Add support for Pick type operator (#64)

Author: hoodmane

type-generation/src/astToIR.ts

@@ -286,6 +286,10 @@ type SyntheticTypeRoot =
   | {
       kind: "omit";
       node: TypeReferenceNode;
+    }
+  | {
+      kind: "pick";
+      node: TypeReferenceNode;
     };
 
 function classifySyntheticType(node: TypeNode): SyntheticTypeRoot | undefined {
@@ -303,11 +307,15 @@ function classifySyntheticType(node: TypeNode): SyntheticTypeRoot | undefined {
   if (name === "Omit") {
     return { kind: "omit", node };
   }
+  if (name === "Pick") {
+    return { kind: "pick", node };
+  }
   return undefined;
 }
 
 type Modifier = {
   omitSet?: Set<string>;
+  pickSet?: Set<string>;
 };
 
 class SyntheticTypeConverter {
@@ -322,14 +330,21 @@ class SyntheticTypeConverter {
   }
 
   hasWork(base: TypeNode, modifiers: Modifier) {
-    const { omitSet } = modifiers;
+    const { omitSet, pickSet } = modifiers;
     if (omitSet) {
       for (const prop of base.getType().getProperties()) {
         if (omitSet.has(prop.getName())) {
           return true;
         }
       }
     }
+    if (pickSet) {
+      for (const prop of base.getType().getProperties()) {
+        if (!pickSet.has(prop.getName())) {
+          return true;
+        }
+      }
+    }
     return false;
   }
 
@@ -342,12 +357,17 @@ class SyntheticTypeConverter {
       name += "_iface";
     }
     let members = nodes.flatMap((x) => x.getMembers());
-    const { omitSet } = modifiers;
+    const { omitSet, pickSet } = modifiers;
     if (omitSet) {
       members = members.filter(
         (x) => !Node.isPropertyNamed(x) || !omitSet.has(x.getName()),
       );
     }
+    if (pickSet) {
+      members = members.filter(
+        (x) => !Node.isPropertyNamed(x) || pickSet.has(x.getName()),
+      );
+    }
     const result = this.converter.interfaceToIR(name, [], members, [], [], []);
     this.converter.extraTopLevels.push(result);
     return { kind: "reference", name, typeArgs: [] };
@@ -395,6 +415,17 @@ class SyntheticTypeConverter {
         this.nameContext.pop();
         return result;
       }
+      case "pick": {
+        const node = typeRoot.node;
+        const base = node.getTypeArguments()[0]!;
+        const toPickType = node.getTypeArguments()[1]!;
+        modifiers = structuredClone(modifiers);
+        modifiers.pickSet = getLiteralTypeArgSet(toPickType);
+        this.nameContext.push("Pick");
+        const res = this.typeToIR(base, modifiers);
+        this.nameContext.pop();
+        return res;
+      }
     }
   }
 

type-generation/tests/a.test.ts

@@ -1597,6 +1597,105 @@ describe("emit", () => {
         );
       });
     });
+    describe("Pick", () => {
+      it("PickLiteral1", () => {
+        const res = emitFile(`
+          type D = Pick<{ a: string; b: number; }, "b">;
+          declare function f(): D;
+        `);
+        assert.strictEqual(
+          removeTypeIgnores(res.slice(1).join("\n\n")),
+          dedent(`
+            type D = D__Pick_iface
+
+            def f() -> D: ...
+
+            class D__Pick_iface(Protocol):
+                b: int | float = ...
+          `).trim(),
+        );
+      });
+      it("PickLiteral2", () => {
+        const res = emitFile(`
+          type D = Pick<{ a: string; b: string; c: number; }, "b" | "c">;
+          declare function f(): D;
+        `);
+        assert.strictEqual(
+          removeTypeIgnores(res.slice(1).join("\n\n")),
+          dedent(`
+            type D = D__Pick_iface
+
+            def f() -> D: ...
+
+            class D__Pick_iface(Protocol):
+                b: str = ...
+                c: int | float = ...
+          `).trim(),
+        );
+      });
+      it("PickIntersection", () => {
+        const res = emitFile(`
+          type D = Pick<{ a: string; b: string; } & { c : string; }, "b">;
+          declare function f(): D;
+        `);
+        assert.strictEqual(
+          removeTypeIgnores(res.slice(1).join("\n\n")),
+          dedent(`
+            type D = D__Pick_iface
+
+            def f() -> D: ...
+
+            class D__Pick__Intersection0_iface(Protocol):
+                b: str = ...
+
+            class D__Pick__Intersection1_iface(Protocol):
+                pass
+
+            class D__Pick_iface(D__Pick__Intersection1_iface, D__Pick__Intersection0_iface, Protocol):
+                pass
+          `).trim(),
+        );
+      });
+      it("PickInterface", () => {
+        const res = emitFile(`
+          interface A { a: string; b: string; }
+          type D = Pick<A, "b">;
+          declare function f(): D;
+        `);
+        assert.strictEqual(
+          removeTypeIgnores(res.slice(1).join("\n\n")),
+          dedent(`
+            type D = D__Pick__A_iface
+
+            def f() -> D: ...
+
+            class D__Pick__A_iface(Protocol):
+                b: str = ...
+          `).trim(),
+        );
+      });
+      it("PickAlias", () => {
+        const res = emitFile(`
+          type A = {
+              s?: boolean;
+              t?: string;
+          };
+          type B = Pick<A, 's'>;
+          declare function f(): B;
+        `);
+        assert.strictEqual(
+          removeTypeIgnores(res.slice(1).join("\n\n")),
+          dedent(`
+            type B = B__Pick_iface
+
+            def f() -> B: ...
+
+            class B__Pick_iface(Protocol):
+                s: bool | None = ...
+          `).trim(),
+        );
+      });
+    });
   });
   describe("adjustments", () => {
     it("setTimeout", () => {