From 3a0682d7172d55816784e952ceef5e42d22a06ee Mon Sep 17 00:00:00 2001
From: Alexander Galabov <galabovalexandar@gmail.com>
Date: Fri, 1 Mar 2024 10:18:03 +0200
Subject: [PATCH] handle enums as keys for zod records

---
 spec/types/record.spec.ts | 83 +++++++++++++++++++++++++++++++++++++++
 src/lib/lodash.ts         |  4 ++
 src/openapi-generator.ts  | 29 +++++++++++++-
 3 files changed, 115 insertions(+), 1 deletion(-)

diff --git a/spec/types/record.spec.ts b/spec/types/record.spec.ts
index a1faea3..0f92c02 100644
--- a/spec/types/record.spec.ts
+++ b/spec/types/record.spec.ts
@@ -59,4 +59,87 @@ describe('record', () => {
       },
     });
   });
+
+  describe('Enum keys', () => {
+    it('supports records with enum keys', () => {
+      const continents = z.enum(['EUROPE', 'AFRICA']);
+
+      const countries = z.enum(['USA', 'CAN']);
+
+      const countryContent = z
+        .object({ countries: countries.array() })
+        .openapi('Content');
+
+      const Geography = z
+        .record(continents, countryContent)
+        .openapi('Geography');
+
+      expectSchema([Geography], {
+        Content: {
+          type: 'object',
+          properties: {
+            countries: {
+              type: 'array',
+              items: {
+                type: 'string',
+                enum: ['USA', 'CAN'],
+              },
+            },
+          },
+          required: ['countries'],
+        },
+
+        Geography: {
+          type: 'object',
+          properties: {
+            EUROPE: { $ref: '#/components/schemas/Content' },
+            AFRICA: { $ref: '#/components/schemas/Content' },
+          },
+        },
+      });
+    });
+
+    it('supports records with native enum keys', () => {
+      enum Continents {
+        EUROPE,
+        AFRICA,
+      }
+
+      const continents = z.nativeEnum(Continents);
+
+      const countries = z.enum(['USA', 'CAN']);
+
+      const countryContent = z
+        .object({ countries: countries.array() })
+        .openapi('Content');
+
+      const Geography = z
+        .record(continents, countryContent)
+        .openapi('Geography');
+
+      expectSchema([Geography], {
+        Content: {
+          type: 'object',
+          properties: {
+            countries: {
+              type: 'array',
+              items: {
+                type: 'string',
+                enum: ['USA', 'CAN'],
+              },
+            },
+          },
+          required: ['countries'],
+        },
+
+        Geography: {
+          type: 'object',
+          properties: {
+            EUROPE: { $ref: '#/components/schemas/Content' },
+            AFRICA: { $ref: '#/components/schemas/Content' },
+          },
+        },
+      });
+    });
+  });
 });
diff --git a/src/lib/lodash.ts b/src/lib/lodash.ts
index ab3d0f0..d6db3e5 100644
--- a/src/lib/lodash.ts
+++ b/src/lib/lodash.ts
@@ -66,3 +66,7 @@ export function uniq<T>(values: T[]) {
 
   return [...set.values()];
 }
+
+export function isString(val: unknown): val is string {
+  return typeof val === 'string';
+}
diff --git a/src/openapi-generator.ts b/src/openapi-generator.ts
index 069c33b..72836e5 100644
--- a/src/openapi-generator.ts
+++ b/src/openapi-generator.ts
@@ -63,6 +63,7 @@ import { enumInfo } from './lib/enum-info';
 import {
   compact,
   isNil,
+  isString,
   mapValues,
   objectEquals,
   omit,
@@ -1009,10 +1010,36 @@ export class OpenAPIGenerator {
 
     if (isZodType(zodSchema, 'ZodRecord')) {
       const propertiesType = zodSchema._def.valueType;
+      const keyType = zodSchema._def.keyType;
+
+      const propertiesSchema = this.generateSchemaWithRef(propertiesType);
+
+      if (
+        isZodType(keyType, 'ZodEnum') ||
+        isZodType(keyType, 'ZodNativeEnum')
+      ) {
+        // Native enums have their keys as both number and strings however the number is an
+        // internal representation and the string is the access point for a documentation
+        const keys = Object.values(keyType.enum).filter(isString);
+
+        const properties = keys.reduce(
+          (acc, curr) => ({
+            ...acc,
+            [curr]: propertiesSchema,
+          }),
+          {} as SchemaObject['properties']
+        );
+
+        return {
+          ...this.mapNullableType('object', isNullable),
+          properties,
+          default: defaultValue,
+        };
+      }
 
       return {
         ...this.mapNullableType('object', isNullable),
-        additionalProperties: this.generateSchemaWithRef(propertiesType),
+        additionalProperties: propertiesSchema,
         default: defaultValue,
       };
     }