Enum primitive for Zod codebases

enum.ts

import * as z from 'zod';

type ReservedFields = 'Schema' | 'KeyAsInputSchema' | 'meta' | '$Key' | '$Type';

type EnumDefinition<K extends string, V extends string | number> = {
  readonly [key in K]: {
    value: V;
    description?: string;
  };
} & {
  readonly [key in ReservedFields]?: never;
};

type EnumMeta<
  Def extends EnumDefinition<string, string | number>,
  K extends keyof Def = keyof Def,
  V extends Def[K]['value'] = Def[K]['value'],
> = {
  definition: Def;

  keys: readonly K[];
  values: readonly V[];
};

export type Enum<
  Def extends EnumDefinition<string, string | number>,
  K extends keyof Def = keyof Def,
  V extends Def[K]['value'] = Def[K]['value'],
> = {
  [Key in K]: Def[Key]['value'] & { $Description: Def[Key]['description'] };
} & {
  meta: EnumMeta<Def>;

  Schema: z.ZodEnum<{ [Key in K]: Def[Key]['value'] }>;
  KeyAsInputSchema: z.ZodPipe<
    z.ZodEnum<{ [Key in K]: Key extends string ? Key : never }>,
    z.ZodTransform<V, K>
  >;

  $Key: K;
  $Type: V;
};

export type $EnumKeys<T> = T extends Enum<infer Def> ? keyof Def : never;
export type $EnumValues<T> = T extends Enum<infer Def>
  ? Def[keyof Def]['value']
  : never;

/**
 * Define an enum
 *
 * @example
 * ```ts
 * const StatusEnum = defineEnum({
 *   Active: { value: 'active' },
 *   Inactive: { value: 'inactive' },
 * });
 * ```
 *
 * @example
 * ```ts
 * StatusEnum.Active; // 'active'
 * StatusEnum.Schema; // z.enum(['active', 'inactive'])
 *
 * typeof StatusEnum.$Key;   // 'Active' | 'Inactive'
 * typeof StatusEnum.$Type; // 'active' | 'inactive'
 * ```
 *
 * Enum values are always normalized as $Value.
 * Use KeyAsInputSchema when using $Key as value.
 *
 * @example
 * ```ts
 * StatusEnum.KeyAsInputSchema.parse('Active'); // 'active'
 * ```
 */
export function defineEnum<const Def extends EnumDefinition<string, string>>(
  definition: Def,
): Enum<Def> {
  const keys: string[] = [];
  const values: string[] = [];
  const enumShape: Record<string, string> = {};
  for (const key in definition) {
    const { value } = definition[key];
    keys.push(key);
    values.push(value);
    enumShape[key] = value;
  }

  const Schema = z.enum(enumShape) as any;
  const KeyAsInputSchema = z.enum(keys).transform(key => enumShape[key]) as any;

  const meta: EnumMeta<Def> = {
    definition,
    keys,
    values,
  };

  return Object.freeze({
    ...enumShape,
    meta,
    Schema,
    KeyAsInputSchema,
    get $Key(): any {
      throw new TypeError(
        "Don't use $Key directly. Use it as a type annotation only.",
      );
    },
    get $Type(): any {
      throw new TypeError(
        "Don't use $Type directly. Use it as a type annotation only.",
      );
    },
  }) as Enum<Def>;
}