fix: handling non-identifier index access type (#289)

* fix: handle single quotes in index access type * fix: handle element access expression in index access type
fabien0102 · Nov 25, 2024 · d9076dd · d9076dd
1 parent fe7e6ae
commit d9076dd
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 5 deletions.
diff --git a/src/core/generateZodSchema.test.ts b/src/core/generateZodSchema.test.ts
@@ -583,6 +583,32 @@ describe("generateZodSchema", () => {
       `);
   });
 
+  it("should deal with index access type using single quote (1st level)", () => {
+    const source = `export type SupermanName = {
+      name?: Superman['name']
+    }`;
+
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "export const supermanNameSchema = z.object({
+          name: supermanSchema.shape.name.optional()
+      });"
+      `);
+  });
+
+  it("should deal with index access type with element access expression (1st level)", () => {
+    const source = `export type SupermanName = {
+      firstName: Superman["name.firstName"]
+      lastName: Superman["name-lastName"]
+    }`;
+
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "export const supermanNameSchema = z.object({
+          firstName: supermanSchema.shape["name.firstName"],
+          lastName: supermanSchema.shape["name-lastName"]
+      });"
+      `);
+  });
+
   it("should deal with record with a union as key", () => {
     const source = `
     export type AvailablePower = Record<Power, boolean>;

diff --git a/src/core/generateZodSchema.ts b/src/core/generateZodSchema.ts
@@ -1390,9 +1390,11 @@ function buildSchemaReference(
     getDependencyName: Required<GenerateZodSchemaProps>["getDependencyName"];
   },
   path = ""
-): ts.PropertyAccessExpression | ts.Identifier {
+): ts.PropertyAccessExpression | ts.Identifier | ts.ElementAccessExpression {
   const indexTypeText = node.indexType.getText(sourceFile);
-  const { indexTypeName, type: indexTypeType } = /^"\w+"$/.exec(indexTypeText)
+  const { indexTypeName, type: indexTypeType } = /^['"]([^'"]+)['"]$/.exec(
+    indexTypeText
+  )
     ? { type: "string" as const, indexTypeName: indexTypeText.slice(1, -1) }
     : { type: "number" as const, indexTypeName: indexTypeText };
 
@@ -1496,9 +1498,19 @@ function buildSchemaReference(
       node.objectType.typeName.getText(sourceFile)
     );
     dependencies.push(dependencyName);
-    return f.createPropertyAccessExpression(
-      f.createIdentifier(dependencyName),
-      f.createIdentifier(`shape.${indexTypeName}.${path}`.slice(0, -1))
+
+    if (/^[a-zA-Z_$][a-zA-Z0-9_$]*$/.test(indexTypeName)) {
+      return f.createPropertyAccessExpression(
+        f.createIdentifier(dependencyName),
+        f.createIdentifier(`shape.${indexTypeName}.${path}`.slice(0, -1))
+      );
+    }
+    return f.createElementAccessExpression(
+      f.createPropertyAccessExpression(
+        f.createIdentifier(dependencyName),
+        f.createIdentifier("shape")
+      ),
+      f.createStringLiteral(indexTypeName)
     );
   }