Support the Pick utility type for a prop (#414)

This PR adds support for the `Pick` TypeScript utility type when defining the type of a prop in a component's prop interface. As mentioned in the previous PR, utility types are still not allowed when defining the overall type of the component's prop interface (i.e. `props: Pick<MyProps, "blah">`). J=SLAP-2968 TEST=auto, manual In the test site, see that using `Pick<ObjectProp, "nestedBool">` instead of `ObjectProp` for the type of `obj` in `BannerData` correctly results in the prop editor for `nestedBool` being the only one that appears in the UI for that prop.
yext · Oct 19, 2023 · b3b57a7 · b3b57a7
1 parent 003f15f
commit b3b57a7
Show file tree

Hide file tree

Showing 11 changed files with 213 additions and 27 deletions.
diff --git a/packages/studio-plugin/src/parsers/PropShapeParser.ts b/packages/studio-plugin/src/parsers/PropShapeParser.ts
@@ -8,7 +8,7 @@ import {
   ParsedShape,
   ParsedType,
   ParsedTypeKind,
-} from "./helpers/TypeNodeParsingHelper";
+} from "./helpers/TypeNodeParsingHelpers";
 
 /**
  * PropShapeParser is a class for parsing a typescript interface into a PropShape.

diff --git a/packages/studio-plugin/src/parsers/StudioSourceFileParser.ts b/packages/studio-plugin/src/parsers/StudioSourceFileParser.ts
@@ -17,7 +17,7 @@ import upath from "upath";
 import vm from "vm";
 import TypeNodeParsingHelper, {
   ParsedType,
-} from "./helpers/TypeNodeParsingHelper";
+} from "./helpers/TypeNodeParsingHelpers";
 import { parseSync as babelParseSync } from "@babel/core";
 import NpmLookup from "./helpers/NpmLookup";
 import { TypelessPropVal } from "../types";

diff --git a/packages/studio-plugin/src/parsers/helpers/StringUnionParsingHelper.ts b/packages/studio-plugin/src/parsers/helpers/StringUnionParsingHelper.ts
@@ -1,5 +1,5 @@
 import { UnionTypeNode, SyntaxKind, LiteralTypeNode, TypeNode } from "ts-morph";
-import { ParsedType, ParsedTypeKind } from "./TypeNodeParsingHelper";
+import { ParsedType, ParsedTypeKind } from "./TypeNodeParsingHelpers";
 
 export default class StringUnionParsingHelper {
   static parseStringUnion(

diff --git a/.../parsers/helpers/TypeNodeParsingHelper.ts → ...parsers/helpers/TypeNodeParsingHelpers.ts b/.../parsers/helpers/TypeNodeParsingHelper.ts → ...parsers/helpers/TypeNodeParsingHelpers.ts
diff --git a/packages/studio-plugin/src/parsers/helpers/UtilityTypeParsingHelper.ts b/packages/studio-plugin/src/parsers/helpers/UtilityTypeParsingHelper.ts
@@ -1,5 +1,14 @@
 import { TypeNode, TypeReferenceNode } from "ts-morph";
-import { ParsedType, ParsedTypeKind } from "./TypeNodeParsingHelper";
+import {
+  ParsedShape,
+  ParsedType,
+  ParsedTypeKind,
+} from "./TypeNodeParsingHelpers";
+
+type UtilityTypeHandler = (
+  typeArgs: TypeNode[],
+  parseTypeNode: (node: TypeNode) => ParsedType
+) => ParsedType;
 
 export default class UtilityTypeParsingHelper {
   static parseUtilityType(
@@ -9,15 +18,18 @@ export default class UtilityTypeParsingHelper {
     const typeName = typeRefNode.getTypeName().getText();
     const typeArgs = typeRefNode.getTypeArguments();
 
-    if (typeName === "Omit") {
-      return this.handleOmit(typeArgs, parseTypeNode);
-    }
+    const utilityTypeHandler: UtilityTypeHandler | undefined = (() => {
+      switch (typeName) {
+        case "Omit":
+          return this.handleOmit;
+        case "Pick":
+          return this.handlePick;
+      }
+    })();
+    return utilityTypeHandler?.(typeArgs, parseTypeNode);
   }
 
-  private static handleOmit(
-    typeArgs: TypeNode[],
-    parseTypeNode: (node: TypeNode) => ParsedType
-  ): ParsedType {
+  private static handleOmit: UtilityTypeHandler = (typeArgs, parseTypeNode) => {
     if (typeArgs.length !== 2) {
       throw new Error(
         `Two type params expected for Omit utility type. Found ${typeArgs.length}.`
@@ -41,5 +53,50 @@ export default class UtilityTypeParsingHelper {
           ` Found ${omitType}.`
       );
     }
-  }
+  };
+
+  private static handlePick: UtilityTypeHandler = (typeArgs, parseTypeNode) => {
+    if (typeArgs.length !== 2) {
+      throw new Error(
+        `Two type params expected for Pick utility type. Found ${typeArgs.length}.`
+      );
+    }
+
+    const [fullType, pickType] = typeArgs.map(parseTypeNode);
+
+    if (fullType.kind !== ParsedTypeKind.Object) {
+      throw new Error(
+        "Only object types are supported for first type arg of Pick utility type." +
+          ` Found ${fullType}.`
+      );
+    }
+    if (
+      !pickType.unionValues &&
+      pickType.kind !== ParsedTypeKind.StringLiteral
+    ) {
+      throw new Error(
+        "Expected string literal or union of string literals for second type arg of Pick utility type." +
+          ` Found ${pickType}.`
+      );
+    }
+
+    const pickKeys = pickType.unionValues ?? [pickType.type];
+    const reducedShape: ParsedShape = pickKeys.reduce((reducedShape, key) => {
+      const parsedProperty = fullType.type[key];
+      if (!parsedProperty) {
+        throw new Error(
+          `Cannot pick key ${key} that is not present in type ${fullType.type}.`
+        );
+      }
+      return {
+        ...reducedShape,
+        [key]: parsedProperty,
+      };
+    }, {});
+
+    return {
+      ...fullType,
+      type: reducedShape,
+    };
+  };
 }
diff --git a/packages/studio-plugin/src/writers/ReactComponentFileWriter.ts b/packages/studio-plugin/src/writers/ReactComponentFileWriter.ts
@@ -25,7 +25,7 @@ import {
 import StudioSourceFileWriter from "./StudioSourceFileWriter";
 import ComponentTreeHelpers from "../utils/ComponentTreeHelpers";
 import camelCase from "camelcase";
-import { CustomTags } from "../parsers/helpers/TypeNodeParsingHelper";
+import { CustomTags } from "../parsers/helpers/TypeNodeParsingHelpers";
 import getImportSpecifier from "../utils/getImportSpecifier";
 
 /**

diff --git a/packages/studio-plugin/tests/parsers/StudioSourceFileParser.test.ts b/packages/studio-plugin/tests/parsers/StudioSourceFileParser.test.ts
@@ -2,7 +2,7 @@ import { SyntaxKind } from "ts-morph";
 import StudioSourceFileParser from "../../src/parsers/StudioSourceFileParser";
 import createTestSourceFile from "../__utils__/createTestSourceFile";
 import expectSyntaxKind from "../__utils__/expectSyntaxKind";
-import { ParsedTypeKind } from "../../src/parsers/helpers/TypeNodeParsingHelper";
+import { ParsedTypeKind } from "../../src/parsers/helpers/TypeNodeParsingHelpers";
 import upath from "upath";
 
 describe("parseExportedObjectLiteral", () => {

diff --git a/...ests/parsers/StaticParsingHelpers.test.ts → ...sers/helpers/StaticParsingHelpers.test.ts b/...ests/parsers/StaticParsingHelpers.test.ts → ...sers/helpers/StaticParsingHelpers.test.ts
@@ -1,13 +1,13 @@
 import { JsxAttributeLike, SyntaxKind } from "ts-morph";
-import { PropShape } from "../../src/types/PropShape";
+import { PropShape } from "../../../src/types/PropShape";
 import {
   PropValueKind,
   PropValueType,
   PropValues,
-} from "../../src/types/PropValues";
-import StaticParsingHelpers from "../../src/parsers/helpers/StaticParsingHelpers";
-import createTestSourceFile from "../__utils__/createTestSourceFile";
-import expectSyntaxKind from "../__utils__/expectSyntaxKind";
+} from "../../../src/types/PropValues";
+import StaticParsingHelpers from "../../../src/parsers/helpers/StaticParsingHelpers";
+import createTestSourceFile from "../../__utils__/createTestSourceFile";
+import expectSyntaxKind from "../../__utils__/expectSyntaxKind";
 
 describe("parseObjectLiteral", () => {
   it("parsing an object literal with an expression", () => {

diff --git a/packages/studio-plugin/tests/parsers/helpers/StringUnionParsingHelper.test.ts b/packages/studio-plugin/tests/parsers/helpers/StringUnionParsingHelper.test.ts
@@ -1,7 +1,7 @@
 import { SyntaxKind } from "ts-morph";
 import StringUnionParsingHelper from "../../../src/parsers/helpers/StringUnionParsingHelper";
 import createTestSourceFile from "../../__utils__/createTestSourceFile";
-import { ParsedTypeKind } from "../../../src/parsers/helpers/TypeNodeParsingHelper";
+import { ParsedTypeKind } from "../../../src/parsers/helpers/TypeNodeParsingHelpers";
 
 it("does not handle StringKeywords within unions", () => {
   const { sourceFile } = createTestSourceFile(

diff --git a/...ers/helpers/TypeNodeParsingHelper.test.ts → ...rs/helpers/TypeNodeParsingHelpers.test.ts b/...ers/helpers/TypeNodeParsingHelper.test.ts → ...rs/helpers/TypeNodeParsingHelpers.test.ts
@@ -1,7 +1,7 @@
 import { SyntaxKind } from "ts-morph";
 import TypeNodeParsingHelper, {
   ParsedTypeKind,
-} from "../../../src/parsers/helpers/TypeNodeParsingHelper";
+} from "../../../src/parsers/helpers/TypeNodeParsingHelpers";
 import createTestSourceFile from "../../__utils__/createTestSourceFile";
 import { PropValueType } from "../../../src/types";
 
@@ -184,7 +184,7 @@ it("throws an error if Array TypeReference is missing a type param", () => {
   ).toThrowError("One type param expected for Array type. Found 0.");
 });
 
-it("can parse a prop with an Omit utility type", () => {
+it("can parse a prop with a utility type", () => {
   const { sourceFile } = createTestSourceFile(
     `export type MyProps = {
       omit: Omit<{ name?: string, nope: number }, "nope">;

diff --git a/packages/studio-plugin/tests/parsers/helpers/UtilityTypeParsingHelper.test.ts b/packages/studio-plugin/tests/parsers/helpers/UtilityTypeParsingHelper.test.ts
@@ -5,7 +5,7 @@ import {
   ParsedProperty,
   ParsedType,
   ParsedTypeKind,
-} from "../../../src/parsers/helpers/TypeNodeParsingHelper";
+} from "../../../src/parsers/helpers/TypeNodeParsingHelpers";
 import cloneDeep from "lodash/cloneDeep";
 
 const nameProp: ParsedProperty = {
@@ -62,10 +62,7 @@ describe("omit", () => {
         type: "num",
       })
     );
-    expect(parsedType?.type).toEqual({
-      name: nameProp,
-      arr: arrProp,
-    });
+    expect(parsedType?.type).toEqual({ name: nameProp, arr: arrProp });
   });
 
   it("can handle multiple omitted types", () => {
@@ -176,3 +173,135 @@ describe("omit", () => {
     );
   });
 });
+
+describe("pick", () => {
+  it("can handle a single picked type", () => {
+    const { sourceFile } = createTestSourceFile(
+      `type MyPick = Pick<{ name?: string, num: number, arr: boolean[] }, "num">;
+      }`
+    );
+    const typeRef = sourceFile.getFirstDescendantByKindOrThrow(
+      SyntaxKind.TypeReference
+    );
+    const parsedType = UtilityTypeParsingHelper.parseUtilityType(
+      typeRef,
+      mockedParseTypeNodeCreator({
+        kind: ParsedTypeKind.StringLiteral,
+        type: "num",
+      })
+    );
+    expect(parsedType?.type).toEqual({ num: numProp });
+  });
+
+  it("can handle multiple picked types", () => {
+    const { sourceFile } = createTestSourceFile(
+      `type MyPick = Pick<{ name?: string, num: number, arr: boolean[] }, "name" | "num">;
+      }`
+    );
+    const typeRef = sourceFile.getFirstDescendantByKindOrThrow(
+      SyntaxKind.TypeReference
+    );
+    const parsedType = UtilityTypeParsingHelper.parseUtilityType(
+      typeRef,
+      mockedParseTypeNodeCreator({
+        kind: ParsedTypeKind.Simple,
+        type: "string",
+        unionValues: ["name", "num"],
+      })
+    );
+    expect(parsedType?.type).toEqual({ name: nameProp, num: numProp });
+  });
+
+  it("throws an error if there an incorrect number of type arguments", () => {
+    const { sourceFile } = createTestSourceFile(
+      `type MyPick = Pick<{ name?: string, num: number, arr: boolean[] }, "name", "num">;
+      }`
+    );
+    const typeRef = sourceFile.getFirstDescendantByKindOrThrow(
+      SyntaxKind.TypeReference
+    );
+
+    expect(() =>
+      UtilityTypeParsingHelper.parseUtilityType(
+        typeRef,
+        mockedParseTypeNodeCreator({
+          kind: ParsedTypeKind.StringLiteral,
+          type: "name",
+        })
+      )
+    ).toThrowError("Two type params expected for Pick utility type. Found 3.");
+  });
+
+  it("throws an error if original type is not an object type", () => {
+    const { sourceFile } = createTestSourceFile(
+      `type MyPick = Pick<"test", "test">;
+      }`
+    );
+    const typeRef = sourceFile.getFirstDescendantByKindOrThrow(
+      SyntaxKind.TypeReference
+    );
+    const testParsedType: ParsedType = {
+      kind: ParsedTypeKind.StringLiteral,
+      type: "test",
+    };
+
+    expect(() =>
+      UtilityTypeParsingHelper.parseUtilityType(
+        typeRef,
+        mockedParseTypeNodeCreator(testParsedType, testParsedType)
+      )
+    ).toThrowError(
+      `Only object types are supported for first type arg of Pick utility type. Found ${testParsedType}.`
+    );
+  });
+
+  it("throws an error if the pick type is not a string literal or union of string literals", () => {
+    const { sourceFile } = createTestSourceFile(
+      `type MyPick = Pick<{ name?: string, num: number, arr: boolean[] }, Array<string>>;
+      }`
+    );
+    const typeRef = sourceFile.getFirstDescendantByKindOrThrow(
+      SyntaxKind.TypeReference
+    );
+
+    const pickType: ParsedType = {
+      kind: ParsedTypeKind.Array,
+      type: {
+        kind: ParsedTypeKind.Simple,
+        type: "string",
+      },
+    };
+
+    expect(() =>
+      UtilityTypeParsingHelper.parseUtilityType(
+        typeRef,
+        mockedParseTypeNodeCreator(pickType)
+      )
+    ).toThrowError(
+      `Expected string literal or union of string literals for second type arg of Pick utility type. Found ${pickType}.`
+    );
+  });
+
+  it("throws an error if pick types include one not a key in original type", () => {
+    const { sourceFile } = createTestSourceFile(
+      `type MyPick = Pick<{ name?: string, num: number, arr: boolean[] }, "name" | "boo">;
+      }`
+    );
+    const typeRef = sourceFile.getFirstDescendantByKindOrThrow(
+      SyntaxKind.TypeReference
+    );
+
+    expect(() =>
+      UtilityTypeParsingHelper.parseUtilityType(
+        typeRef,
+        mockedParseTypeNodeCreator({
+          kind: ParsedTypeKind.Simple,
+          type: "string",
+          unionValues: ["name", "boo"],
+        })
+      )
+    ).toThrowError(
+      `Cannot pick key boo that is not present in type ${fullObjectType.type}.`
+    );
+  });
+});