Support arbitrary keys in objects by falling back to the functional typed dict definition (#21)

* update gitignore for mac

* use the workspace ts version

* support arbitrary keys in dicts by falling back to a functional typed dict definition

* improve readability of functional dicts

Author: TheLartians

.gitignore

@@ -2,3 +2,4 @@
 /build
 /yarn-error.log
 /dist
+.DS_Store

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "js/ts.tsdk.path": "node_modules/typescript/lib"
+}

src/parseProperty.ts

@@ -34,3 +34,42 @@ export const parseProperty = (state: ParserState, symbol: ts.Symbol) => {
     return `${name}: ${definition}${documentationSuffix}`;
   }
 };
+
+export const parsePropertyForDict = (state: ParserState, symbol: ts.Symbol) => {
+  const name = JSON.stringify(symbol.getName());
+  if (symbol.flags & ts.SymbolFlags.Optional) {
+      state.imports.add("NotRequired");
+      const definition = parseInlineType(
+        state,
+        // since the entry is already options, the inner type can be non-nullable
+        state.typechecker.getNonNullableType(state.typechecker.getTypeOfSymbol(symbol)),
+      );
+    if (state.config.nullableOptionals) {
+      state.imports.add("Optional");
+      return `${name}: NotRequired[Optional[${definition}]]`;
+    } else {
+      return `${name}: NotRequired[${definition}]`;
+    }
+  } else {
+    const definition = parseInlineType(
+      state,
+      // since the entry is already options, the inner type can be non-nullable
+      state.typechecker.getTypeOfSymbol(symbol),
+    );
+    return `${name}: ${definition}`;
+  }
+};
+
+export const getDocumentationStringForDict = (state: ParserState, symbol: ts.Symbol) => {
+  const name = symbol.getName();
+  const documentation = symbol
+    .getDocumentationComment(state.typechecker)
+    .map((v) => v.text)
+    .join("  \n");
+  if (documentation.length > 0) {
+    return `${JSON.stringify(name)}: ${documentation}`
+  } else {
+    return undefined;
+  }
+} 
+

src/parseTypeDefinition.ts

@@ -1,6 +1,6 @@
 import ts from "typescript";
 import { ParserState } from "./ParserState";
-import { parseProperty } from "./parseProperty";
+import { getDocumentationStringForDict, parseProperty, parsePropertyForDict } from "./parseProperty";
 import { getDocumentationStringForType } from "./getDocumentationStringForType";
 import { tryToParseInlineType } from "./parseInlineType";
 import { isValidPythonIdentifier } from "./isValidPythonIdentifier";
@@ -25,18 +25,40 @@ export const parseTypeDefinition = (
     state.statements.push(definition);
   } else {
     state.imports.add("TypedDict");
-
-    const properties = type
+    
+    const allKeysAreValidPythonIdentifiers = type
       .getProperties()
-      .filter((v) => isValidPythonIdentifier(v.getName()))
-      .map((v) => parseProperty(state, v));
+      .map((v) => isValidPythonIdentifier(v.getName()))
+      .reduce((a,b) => a&&b, true);
 
-      const definition = `class ${name}(TypedDict):${
-      documentation
-        ? `\n  """\n  ${documentation.replaceAll("\n", "  \n")}\n  """`
-        : ""
-    }\n  ${properties.length > 0 ? properties.join(`\n  `) : "pass"}`;
+    if (allKeysAreValidPythonIdentifiers) {
+      const properties = type
+        .getProperties()
+        .map((v) => parseProperty(state, v));
+  
+        const definition = `class ${name}(TypedDict):${
+        documentation
+          ? `\n  """\n  ${documentation.replaceAll("\n", "  \n")}\n  """`
+          : ""
+      }\n  ${properties.length > 0 ? properties.join(`\n  `) : "pass"}`;
+      state.statements.push(definition);
+    } else {
+      const properties = type
+        .getProperties()
+        .map((v) => parsePropertyForDict(state, v));
+
+      const propertyDocumentation = type
+        .getProperties()
+        .map((v) => getDocumentationStringForDict(state, v))
+        .filter(v => v !== undefined)
+        .join("\n");
+
+      const innerDocstring = documentation?.replaceAll("\n", "  \n") + (propertyDocumentation.length > 0 ? "\n## Entries\n" + propertyDocumentation : "");
+      const docstring = innerDocstring.length > 0 ? `\n"""\n${innerDocstring}\n"""` : "";
+      const definition = `${name} = TypedDict(${JSON.stringify(name)}, {\n  ${properties.join(",\n  ")}\n})${docstring}`;
+
+      state.statements.push(definition);
+    }
 
-    state.statements.push(definition);
   }
 };

src/testing/dicts.test.ts

@@ -59,12 +59,12 @@ class A(TypedDict):
     expect(result).toContain(`class A(TypedDict):\n  foo: str\n  bar: float`);
   });
 
-  it("transpiles optional values as NotRequired[Optional[T]]", async () => {
+  it("transpiles optional values as NotRequired[T]", async () => {
     const result = await transpileString(`export type A = { foo?: string }`);
     expect(result).toContain(`class A(TypedDict):\n  foo: NotRequired[str]`);
   });
 
-  it("transpiles optional values as NotRequired[Optional[T]] in strict mode", async () => {
+  it("transpiles optional values as NotRequired[T] in strict mode", async () => {
     const result = await transpileString(
       `export type A = { foo?: string }`,
       {},
@@ -73,7 +73,7 @@ class A(TypedDict):
     expect(result).toContain(`class A(TypedDict):\n  foo: NotRequired[str]`);
   });
 
-  it("transpiles optional values with non-null optionals as NotRequired[T]", async () => {
+  it("transpiles optional values with non-null optionals as NotRequired[Optional[T]]", async () => {
     const result = await transpileString(`export type A = { foo?: string }`, {
       nullableOptionals: true,
     });
@@ -88,4 +88,39 @@ class A(TypedDict):
     );
     expect(result).toContain(`class A(TypedDict):\n  foo: float\n  bar: float`);
   });
+
+  it("falls back to the functional syntax if keys are unsupported", async () => {
+    const result = await transpileString(`export type A = {
+      "foo.bar"?: string,
+    }`);
+    expect(result).toContain(
+      `A = TypedDict("A", {
+  "foo.bar": NotRequired[str]
+})`,
+    );
+  });
+
+  it("moves the key/value docstrings to the object docstring in the functional syntax", async () => {
+    const result = await transpileString(`
+      /** This is A */
+      export type A = {
+        /** this is foo.bar */
+        "foo.bar": string,
+        /** this is a/b */
+        "a/b": number,
+        "undocumented": string,
+      }
+    `);
+    expect(result).toContain(`A = TypedDict("A", {
+  "foo.bar": str,
+  "a/b": float,
+  "undocumented": str
+})
+"""
+This is A
+## Entries
+"foo.bar": this is foo.bar
+"a/b": this is a/b
+"""`);
+  });
 });