feat: Improved Promise handling to support packages like Prisma (#1924)

arthurfiorette · web-flow · commit 183b4264f554 · 2024-05-20T11:15:21.000-10:00
diff --git a/.gitignore b/.gitignore
@@ -8,3 +8,6 @@ node_modules/
 # local config for auto
 .env
 
+# Other package managers
+pnpm-lock.yaml
+package-lock.json
diff --git a/README.md b/README.md
@@ -231,6 +231,7 @@ fs.writeFile(outputPath, schemaString, (err) => {
 -   `keyof`
 -   conditional types
 -   functions
+-   `Promise<T>` unwraps to `T`
 
 ## Run locally
 
diff --git a/factory/parser.ts b/factory/parser.ts
@@ -1,12 +1,12 @@
-import ts from "typescript";
+import type ts from "typescript";
 import { BasicAnnotationsReader } from "../src/AnnotationsReader/BasicAnnotationsReader.js";
 import { ExtendedAnnotationsReader } from "../src/AnnotationsReader/ExtendedAnnotationsReader.js";
 import { ChainNodeParser } from "../src/ChainNodeParser.js";
 import { CircularReferenceNodeParser } from "../src/CircularReferenceNodeParser.js";
-import { CompletedConfig } from "../src/Config.js";
+import type { CompletedConfig } from "../src/Config.js";
 import { ExposeNodeParser } from "../src/ExposeNodeParser.js";
-import { MutableParser } from "../src/MutableParser.js";
-import { NodeParser } from "../src/NodeParser.js";
+import type { MutableParser } from "../src/MutableParser.js";
+import type { NodeParser } from "../src/NodeParser.js";
 import { AnnotatedNodeParser } from "../src/NodeParser/AnnotatedNodeParser.js";
 import { AnyTypeNodeParser } from "../src/NodeParser/AnyTypeNodeParser.js";
 import { ArrayLiteralExpressionNodeParser } from "../src/NodeParser/ArrayLiteralExpressionNodeParser.js";
@@ -55,9 +55,10 @@ import { UndefinedTypeNodeParser } from "../src/NodeParser/UndefinedTypeNodePars
 import { UnionNodeParser } from "../src/NodeParser/UnionNodeParser.js";
 import { UnknownTypeNodeParser } from "../src/NodeParser/UnknownTypeNodeParser.js";
 import { VoidTypeNodeParser } from "../src/NodeParser/VoidTypeNodeParser.js";
-import { SubNodeParser } from "../src/SubNodeParser.js";
+import type { SubNodeParser } from "../src/SubNodeParser.js";
 import { TopRefNodeParser } from "../src/TopRefNodeParser.js";
 import { SatisfiesNodeParser } from "../src/NodeParser/SatisfiesNodeParser.js";
+import { PromiseNodeParser } from "../src/NodeParser/PromiseNodeParser.js";
 
 export type ParserAugmentor = (parser: MutableParser) => void;
 
@@ -121,6 +122,7 @@ export function createParser(program: ts.Program, config: CompletedConfig, augme
         .addNodeParser(new LiteralNodeParser(chainNodeParser))
         .addNodeParser(new ParenthesizedNodeParser(chainNodeParser))
 
+        .addNodeParser(new PromiseNodeParser(typeChecker, chainNodeParser))
         .addNodeParser(new TypeReferenceNodeParser(typeChecker, chainNodeParser))
         .addNodeParser(new ExpressionWithTypeArgumentsNodeParser(typeChecker, chainNodeParser))
         .addNodeParser(new IndexedAccessTypeNodeParser(typeChecker, chainNodeParser))
diff --git a/package.json b/package.json
@@ -65,6 +65,7 @@
     "@types/jest": "^29.5.12",
     "@types/node": "^20.12.7",
     "@types/normalize-path": "^3.0.2",
+    "@types/ts-expose-internals": "npm:ts-expose-internals@^5.4.5",
     "ajv": "^8.12.0",
     "ajv-formats": "^3.0.1",
     "auto": "^11.1.6",
@@ -94,6 +95,5 @@
     "debug": "tsx --inspect-brk ts-json-schema-generator.ts",
     "run": "tsx ts-json-schema-generator.ts",
     "release": "yarn build && auto shipit"
-  },
-  "packageManager": "yarn@1.22.19"
+  }
 }
diff --git a/src/NodeParser/FunctionNodeParser.ts b/src/NodeParser/FunctionNodeParser.ts
@@ -18,8 +18,11 @@ export class FunctionNodeParser implements SubNodeParser {
     public supportsNode(node: ts.TypeNode): boolean {
         return (
             node.kind === ts.SyntaxKind.FunctionType ||
+            // @ts-expect-error internals type bug
             node.kind === ts.SyntaxKind.FunctionExpression ||
+            // @ts-expect-error internals type bug
             node.kind === ts.SyntaxKind.ArrowFunction ||
+            // @ts-expect-error internals type bug
             node.kind === ts.SyntaxKind.FunctionDeclaration
         );
     }
diff --git a/src/NodeParser/PromiseNodeParser.ts b/src/NodeParser/PromiseNodeParser.ts
@@ -0,0 +1,98 @@
+import ts from "typescript";
+import { Context, type NodeParser } from "../NodeParser.js";
+import type { SubNodeParser } from "../SubNodeParser.js";
+import { AliasType } from "../Type/AliasType.js";
+import type { BaseType } from "../Type/BaseType.js";
+import { DefinitionType } from "../Type/DefinitionType.js";
+import { getKey } from "../Utils/nodeKey.js";
+
+/**
+ * Needs to be registered before 261, 260, 230, 262 node kinds
+ */
+export class PromiseNodeParser implements SubNodeParser {
+    public constructor(
+        protected typeChecker: ts.TypeChecker,
+        protected childNodeParser: NodeParser,
+    ) {}
+
+    public supportsNode(node: ts.Node): boolean {
+        if (
+            // 261  interface PromiseInterface extends Promise<T>
+            !ts.isInterfaceDeclaration(node) &&
+            // 260  class PromiseClass implements Promise<T>
+            !ts.isClassDeclaration(node) &&
+            // 230  Promise<T>
+            !ts.isExpressionWithTypeArguments(node) &&
+            // 262  type PromiseAlias = Promise<T>;
+            !ts.isTypeAliasDeclaration(node)
+        ) {
+            return false;
+        }
+
+        const type = this.typeChecker.getTypeAtLocation(node);
+
+        const awaitedType = this.typeChecker.getAwaitedType(type);
+
+        // ignores non awaitable types
+        if (!awaitedType) {
+            return false;
+        }
+
+        // If the awaited type differs from the original type, the type extends promise
+        // Awaited<Promise<T>> -> T (Promise<T> !== T)
+        // Awaited<Y> -> Y (Y === Y)
+        if (awaitedType === type) {
+            return false;
+        }
+
+        // In types like: A<T> = T, type C = A<1>, C has the same type as A<1> and 1,
+        // the awaitedType is NOT the same reference as the type, so a assignability
+        // check is needed
+        return (
+            !this.typeChecker.isTypeAssignableTo(type, awaitedType) &&
+            !this.typeChecker.isTypeAssignableTo(awaitedType, type)
+        );
+    }
+
+    public createType(
+        node: ts.InterfaceDeclaration | ts.ClassDeclaration | ts.ExpressionWithTypeArguments | ts.TypeAliasDeclaration,
+        context: Context,
+    ): BaseType {
+        const type = this.typeChecker.getTypeAtLocation(node);
+        const awaitedType = this.typeChecker.getAwaitedType(type)!; // supportsNode ensures this
+        const awaitedNode = this.typeChecker.typeToTypeNode(awaitedType, undefined, ts.NodeBuilderFlags.IgnoreErrors);
+
+        if (!awaitedNode) {
+            throw new Error(
+                `Could not find awaited node for type ${node.pos === -1 ? "<unresolved>" : node.getText()}`,
+            );
+        }
+
+        const baseNode = this.childNodeParser.createType(awaitedNode, new Context(node));
+
+        const name = this.getNodeName(node);
+
+        // Nodes without name should just be their awaited type
+        // export class extends Promise<T> {} -> T
+        // export class A extends Promise<T> {} -> A (ref to T)
+        if (!name) {
+            return baseNode;
+        }
+
+        return new DefinitionType(name, new AliasType(`promise-${getKey(node, context)}`, baseNode));
+    }
+
+    private getNodeName(
+        node: ts.InterfaceDeclaration | ts.ClassDeclaration | ts.ExpressionWithTypeArguments | ts.TypeAliasDeclaration,
+    ) {
+        if (ts.isExpressionWithTypeArguments(node)) {
+            if (!ts.isHeritageClause(node.parent)) {
+                throw new Error("Expected ExpressionWithTypeArguments to have a HeritageClause parent");
+            }
+
+            return node.parent.parent.name?.getText();
+        }
+
+        return node.name?.getText();
+    }
+}
diff --git a/src/NodeParser/TypeReferenceNodeParser.ts b/src/NodeParser/TypeReferenceNodeParser.ts
@@ -1,6 +1,5 @@
 import ts from "typescript";
-
-import { Context, NodeParser } from "../NodeParser.js";
+import { Context, type NodeParser } from "../NodeParser.js";
 import type { SubNodeParser } from "../SubNodeParser.js";
 import { AnnotatedType } from "../Type/AnnotatedType.js";
 import { AnyType } from "../Type/AnyType.js";
@@ -31,11 +30,6 @@ export class TypeReferenceNodeParser implements SubNodeParser {
             // property on the node itself.
             (node.typeName as unknown as ts.Type).symbol;
 
-        // Wraps promise type to avoid resolving to a empty Object type.
-        if (typeSymbol.name === "Promise") {
-            return this.childNodeParser.createType(node.typeArguments![0], this.createSubContext(node, context));
-        }
-
         if (typeSymbol.flags & ts.SymbolFlags.Alias) {
             const aliasedSymbol = this.typeChecker.getAliasedSymbol(typeSymbol);
 
@@ -53,6 +47,16 @@ export class TypeReferenceNodeParser implements SubNodeParser {
             return context.getArgument(typeSymbol.name);
         }
 
+        // Wraps promise type to avoid resolving to a empty Object type.
+        if (typeSymbol.name === "Promise" || typeSymbol.name === "PromiseLike") {
+            // Promise without type resolves to Promise<any>
+            if (!node.typeArguments || node.typeArguments.length === 0) {
+                return new AnyType();
+            }
+
+            return this.childNodeParser.createType(node.typeArguments[0], this.createSubContext(node, context));
+        }
+
         if (typeSymbol.name === "Array" || typeSymbol.name === "ReadonlyArray") {
             const type = this.createSubContext(node, context).getArguments()[0];
 
diff --git a/test/utils.ts b/test/utils.ts
@@ -69,7 +69,7 @@ export function assertValidSchema(
         const actual: any = JSON.parse(JSON.stringify(schema));
 
         expect(typeof actual).toBe("object");
-        expect(actual).toEqual(expected);
+        expect(actual).toStrictEqual(expected);
 
         let localValidator = validator;
         if (config.extraTags) {
diff --git a/test/valid-data-type.test.ts b/test/valid-data-type.test.ts
@@ -143,4 +143,6 @@ describe("valid-data-type", () => {
     it("ignore-export", assertValidSchema("ignore-export", "*"));
 
     it("lowercase", assertValidSchema("lowercase", "MyType"));
+
+    it("promise-extensions", assertValidSchema("promise-extensions", "*"));
 });
diff --git a/test/valid-data/promise-extensions/main.ts b/test/valid-data/promise-extensions/main.ts
@@ -0,0 +1,59 @@
+export type A = { a: string; b: number[] };
+
+export type PromiseAlias = Promise<A>;
+
+export class PromiseClass extends Promise<A> {}
+
+export interface PromiseInterface extends Promise<A> {}
+
+export type LikeType = PromiseLike<A>;
+
+export type PromiseOrAlias = Promise<A> | A;
+
+export type LikeOrType = PromiseLike<A> | A;
+
+export type AndPromise = Promise<A> & { a: string };
+
+export type AndLikePromise = PromiseLike<A> & { a: string };
+
+// Should not be present
+export default class extends Promise<A> {}
+
+export class LikeClass implements PromiseLike<A> {
+    then<TResult1 = A, TResult2 = never>(
+        onfulfilled?: ((value: A) => TResult1 | PromiseLike<TResult1>) | null | undefined,
+        onrejected?: ((reason: unknown) => TResult2 | PromiseLike<TResult2>) | null | undefined
+    ): PromiseLike<TResult1 | TResult2> {
+        return new Promise(() => {});
+    }
+}
+
+export abstract class LikeAbstractClass implements PromiseLike<A> {
+    abstract then<TResult1 = A, TResult2 = never>(
+        onfulfilled?: ((value: A) => TResult1 | PromiseLike<TResult1>) | null | undefined,
+        onrejected?: ((reason: unknown) => TResult2 | PromiseLike<TResult2>) | null | undefined
+    );
+}
+
+export interface LikeInterface extends PromiseLike<A> {}
+
+// Prisma has a base promise type just like this
+export interface WithProperty extends Promise<A> {
+    [Symbol.toStringTag]: "WithProperty";
+}
+
+export interface ThenableInterface {
+    then<TResult1 = A, TResult2 = never>(
+        onfulfilled?: ((value: A) => TResult1 | PromiseLike<TResult1>) | null | undefined,
+        onrejected?: ((reason: unknown) => TResult2 | PromiseLike<TResult2>) | null | undefined
+    ): PromiseLike<TResult1 | TResult2>;
+}
+
+export class ThenableClass {
+    then<TResult1 = A, TResult2 = never>(
+        onfulfilled?: ((value: A) => TResult1 | PromiseLike<TResult1>) | null | undefined,
+        onrejected?: ((reason: unknown) => TResult2 | PromiseLike<TResult2>) | null | undefined
+    ): PromiseLike<TResult1 | TResult2> {
+        return new Promise(() => {});
+    }
+}
diff --git a/test/valid-data/promise-extensions/schema.json b/test/valid-data/promise-extensions/schema.json
@@ -0,0 +1,66 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "definitions": {
+    "A": {
+      "additionalProperties": false,
+      "properties": {
+        "a": {
+          "type": "string"
+        },
+        "b": {
+          "items": {
+            "type": "number"
+          },
+          "type": "array"
+        }
+      },
+      "required": [
+        "a",
+        "b"
+      ],
+      "type": "object"
+    },
+    "AndLikePromise": {
+      "$ref": "#/definitions/A"
+    },
+    "AndPromise": {
+      "$ref": "#/definitions/A"
+    },
+    "LikeAbstractClass": {
+      "$ref": "#/definitions/A"
+    },
+    "LikeClass": {
+      "$ref": "#/definitions/A"
+    },
+    "LikeInterface": {
+      "$ref": "#/definitions/A"
+    },
+    "LikeOrType": {
+      "$ref": "#/definitions/A"
+    },
+    "LikeType": {
+      "$ref": "#/definitions/A"
+    },
+    "PromiseAlias": {
+      "$ref": "#/definitions/A"
+    },
+    "PromiseClass": {
+      "$ref": "#/definitions/A"
+    },
+    "PromiseInterface": {
+      "$ref": "#/definitions/A"
+    },
+    "PromiseOrAlias": {
+      "$ref": "#/definitions/A"
+    },
+    "ThenableClass": {
+      "$ref": "#/definitions/A"
+    },
+    "ThenableInterface": {
+      "$ref": "#/definitions/A"
+    },
+    "WithProperty": {
+      "$ref": "#/definitions/A"
+    }
+  }
+}
diff --git a/test/vega-lite/schema.json b/test/vega-lite/schema.json
@@ -27149,6 +27149,8 @@
     },
     "SingleDefUnitChannel": {
       "enum": [
+        "text",
+        "shape",
         "x",
         "y",
         "xOffset",
@@ -27173,9 +27175,7 @@
         "strokeDash",
         "size",
         "angle",
-        "shape",
         "key",
-        "text",
         "href",
         "url",
         "description"
diff --git a/yarn.lock b/yarn.lock
@@ -1976,6 +1976,11 @@
   resolved "https://registry.yarnpkg.com/@types/stack-utils/-/stack-utils-2.0.3.tgz#6209321eb2c1712a7e7466422b8cb1fc0d9dd5d8"
   integrity sha512-9aEbYZ3TbYMznPdcdr3SmIrLXwC/AKZXQeCf9Pgao5CKb8CyHuEX5jzWPTkvregvhRJHcpRO6BFoGW9ycaOkYw==
 
+"@types/ts-expose-internals@npm:ts-expose-internals@^5.4.5":
+  version "5.4.5"
+  resolved "https://registry.yarnpkg.com/ts-expose-internals/-/ts-expose-internals-5.4.5.tgz#94da2b665627135ad1281d98af3ccb08cb4c1950"
+  integrity sha512-0HfRwjgSIOyuDlHzkFedMWU4aHWq9pu4MUKHgH75U+L76wCAtK5WB0rc/dAIhulMRcPUlcKONeiiR5Sxy/7XcA==
+
 "@types/yargs-parser@*":
   version "21.0.3"
   resolved "https://registry.yarnpkg.com/@types/yargs-parser/-/yargs-parser-21.0.3.tgz#815e30b786d2e8f0dcd85fd5bcf5e1a04d008f15"
