zod cli parser

a_readme.md

example defineCommand

import { z } from 'zod';

import { defineCommand } from './src/define-command.ts';
import { defineOptionsRaw } from './src/define-options.ts';

export const lint = defineCommand({
  name: 'lint',
  description: 'Run lint',
  args: [z.string().min(5), z.coerce.number().default(123)],

  // ?NOTE: `aliases` is array of command aliases which can be used to run the command
  aliases: ['li', 'lnt', 'lnit'],

  // ?NOTE: inside `defineCommand`, we should manually call one of the methods returned from `defineOptions`,
  // (which are the `run`, `parse`, `parseAsync`, `safeParse`, and `safeParseAsync`),
  // using all the process.argv (`argv`) up until the first non-flag argument,
  // eg. the "command name" (in this case, `lint`).
  options: defineOptionsRaw({
    schema: z.object({
      fix: z.boolean().default(true),
      quiet: z.boolean(),
    }),
    aliases: {
      f: 'fix',
      q: 'quiet',
    },
  }),

  // ?NOTE: where `action` is the function that will be called when the command is run.
  // It receives the parsed `options` from above as `commandOpptions` preserving the types.
  // It also receives the `args` is an array of types defined in the above `args` definition.
  action: (commandOptions, args) => {
    console.log(
      'bruh',
      commandOptions,
      args,
      // { args: { name, num } },
      // name.endsWith('foo'),
      // num.toFixed(2),
    );
  },
});

// console.log(lint);

const res = lint.parse();

console.log('result:', res);

example defineOptions

import { z } from 'zod';

import { defineOptions } from './src/define-options.ts';

export const globalOptions = defineOptions({
  schema: z.object({
    replicas: z.coerce.number().int().min(1).max(10),
    tags: z.array(z.string()).min(1),
    env: z.record(z.string(), z.string()),
    timeout: z.coerce.number().positive().default(30),
    role: z.enum(['admin', 'user']),
    active: z.boolean(),
    barry: z.boolean(),
  }),
  aliases: {
    r: 'role',
    a: 'active',
    t: 'timeout',
    e: 'timeout',
    p: 'replicas',
    g: 'tags',
    b: 'barry',
    // foo: 'missing', // it's EXPECTED to ERROR here, because "missing" is not a key in the `options` schema
    // role: 'timeout', // it's EXPECTED to ERROR here, because `role` is a key in the `options` schema
  },
});

const result = globalOptions.parse(
  // ['--foo=bar', '-r=user', '-ep=5', '-g=foo', '--tags=barry', '-g=rar', '--tags=woo', '-ab'],
  // { allowUnknown: true },
  ['-r=user', '-ep=5', '-g=foo', '--tags=barry', '-g=rar', '--tags=woo', '-ab'],
  // { allowUnknown: false }, // default: false
);

console.log('bruh', result);

define-command.ts

import { z } from 'zod';

import { defineOptions } from './define-options.ts';
import type {
  DefineCommandConfig,
  DefineOptionsConfig,
  ParseConfig,
  ResultData,
  SafeFailure,
  SafeSuccess,
  StandardDefineReturn,
} from './types.ts';

export function defineCommandRaw<T extends z.ZodObject<z.ZodRawShape>>(
  config: DefineCommandConfig<T>,
): DefineCommandConfig<T> {
  return config;
}

export function defineCommand<
  T extends z.ZodObject<z.ZodRawShape>,
  CommandArgs extends z.ZodTypeAny[] = z.ZodTypeAny[],
>({
  name,
  // description,
  args,
  aliases,
  options,
  action,
}: DefineCommandConfig<T>): StandardDefineReturn<T> {
  const issues: any[] = [];
  const ctx = { addIssue: (issue: any) => issues.push(issue) };

  return {
    run,
    parse,
    parseAsync,
    safeParse,
    safeParseAsync,
  };

  // TODO 1: use `allowUnknown`, maybe allow unknown commands to call the `action` function
  // TODO 2: add runAsync, to call `action` with await; use `runAsync` in `*parseAsync` methods
  function run(argv?: string[] | null, config?: ParseConfig) {
    const cfg = { allowUnknown: false, safe: false, ...config };
    argv = argv || process.argv.slice(2);

    const commandName = argv.shift() || '';
    const noCommand = commandName !== name && !aliases?.includes(commandName);

    if (noCommand) {
      ctx.addIssue({
        code: z.ZodIssueCode.custom,
        message: `No such command: ${commandName}`,
        path: [commandName],
      });

      if (!cfg.safe) {
        throw new z.ZodError(issues);
      }

      return { success: false, error: new z.ZodError(issues) } as SafeFailure<T>;
    }

    const commandIndex = argv.findIndex((arg) => arg.startsWith('-'));
    const commandArgs = (commandIndex === -1 ? argv : argv?.slice(0, commandIndex + 1)) || [];

    // Parse options
    const cmdOptions = defineOptions(
      (options || { schema: z.object({}) }) as DefineOptionsConfig<T>,
    );
    const parseMethod = cfg.safe ? cmdOptions.safeParse : cmdOptions.parse;
    const cmdOptsResult = parseMethod(argv.slice(0, commandIndex), config);

    // Validate and parse args
    const parsedArgs =
      args?.map((argSchema, index) => {
        const argValue = commandArgs[index];

        // TODO: add support for `safeParse`
        return argSchema.parse(argValue);
      }) || [];

    const cmdArgs = parsedArgs as { [K in keyof CommandArgs]: z.infer<CommandArgs[K]> };

    if (cfg.safe) {
      const optResult = cmdOptsResult as SafeSuccess<T> | SafeFailure<T>;

      if (optResult.success) {
        // ?NOTE: TypeScript tricks for types
        const opts = optResult.data.options as z.infer<typeof optResult.data.schema>;

        // TODO: support returning this `result` along the `options` and `schema`
        const _result = action(opts, cmdArgs);

        return {
          success: true,
          data: { options: opts, schema: optResult.data.schema },
        } as SafeSuccess<T>;
      }

      // ?NOTE: TypeScript tricks for types
      const werrSchema = optResult.error.schema;
      (optResult.error as any).schema = werrSchema;

      return { success: false, error: optResult.error } as SafeFailure<T>;
    }

    // return { success: false, error: { message: 'not implemented' } };
    const optResult = cmdOptsResult as ResultData<T>;
    const opts = optResult.options;

    // TODO: support returning this `result` along the `options` and `schema`
    const _result = action(opts, cmdArgs);

    return { options: opts, schema: optResult.schema } as ResultData<T>;
  }

  function parse(argv?: string[] | null, config?: ParseConfig): ResultData<T> {
    return run(argv, { ...config, safe: false }) as ResultData<T>;
  }

  async function parseAsync(argv?: string[] | null, config?: ParseConfig): Promise<ResultData<T>> {
    return run(argv, { ...config, safe: false }) as ResultData<T>;
  }

  function safeParse(
    argv?: string[] | null,
    config?: ParseConfig,
  ): SafeSuccess<T> | SafeFailure<T> {
    return run(argv, { ...config, safe: true }) as SafeSuccess<T> | SafeFailure<T>;
  }

  async function safeParseAsync(
    argv?: string[] | null,
    config?: ParseConfig,
  ): Promise<SafeSuccess<T> | SafeFailure<T>> {
    return run(argv, { ...config, safe: true }) as SafeSuccess<T> | SafeFailure<T>;
  }
}

define-options.ts

import { z } from 'zod';

import type {
  DefineOptionsConfig,
  ParseConfig,
  ResultData,
  SafeFailure,
  SafeSuccess,
  StandardDefineReturn,
} from './types.ts';
import { parseFlags, transformResults } from './utils.ts';

export function defineOptionsRaw<T extends z.ZodObject<z.ZodRawShape>>(
  config: DefineOptionsConfig<T>,
): DefineOptionsConfig<T> {
  return config;
}

export function defineOptions<T extends z.ZodObject<z.ZodRawShape>>({
  schema: optionsSchema,
  aliases,
}: DefineOptionsConfig<T>): StandardDefineReturn<T> {
  const issues: any[] = [];
  const ctx = { addIssue: (issue: any) => issues.push(issue) };

  return {
    run,
    parse,
    parseAsync,
    safeParse,
    safeParseAsync,
  };

  function run(argv?: string[] | null, config?: ParseConfig) {
    const cfg = { allowUnknown: false, safe: false, ...config };
    argv = argv || process.argv.slice(2);

    // Parse global flags
    //
    // { schema, aliases }: DefineOptionsConfig<T>,
    // { allowUnknown }: ParseConfig = { allowUnknown: false },
    const { rawResults: globalResults, unknownKeys } = parseFlags(
      argv,
      { schema: optionsSchema, aliases },
      cfg,
    );

    if (unknownKeys.length > 0) {
      ctx.addIssue({
        code: z.ZodIssueCode.unrecognized_keys,
        message: `Unrecognized key(s) in options: ${unknownKeys.join(', ')}`,
        keys: unknownKeys,
      });
      if (!cfg.safe) {
        throw new z.ZodError(issues);
      }
    }

    for (const [opt, optSchema] of Object.entries(optionsSchema.shape)) {
      if (optSchema._def.typeName !== 'ZodDefault') {
        optionsSchema.shape[opt] = optSchema.optional();
      }
    }

    const schema = cfg.allowUnknown ? optionsSchema.passthrough() : optionsSchema;
    const parseMethod = cfg.safe ? schema.safeParse : schema.parse;
    // Parse and validate global options
    // ?NOTE: it doesn't merge the aliases, not needed, they are only for the command-line
    const options = parseMethod(transformResults(globalResults));

    if (cfg.safe) {
      if (options.success) {
        return {
          success: true,
          data: { schema, options: options.data },
        } as SafeSuccess<T>;
      }

      options.error.schema = schema;
      return { success: false, error: options.error } as {
        success: boolean;
        error: typeof options.error;
      } as SafeFailure<T>;
    }

    return { schema, options } as ResultData<T>;
  }

  function parse(argv?: string[] | null, config?: ParseConfig): ResultData<T> {
    return run(argv, { ...config, safe: false }) as ResultData<T>;
  }

  async function parseAsync(argv?: string[] | null, config?: ParseConfig): Promise<ResultData<T>> {
    return run(argv, { ...config, safe: false }) as ResultData<T>;
  }

  function safeParse(
    argv?: string[] | null,
    config?: ParseConfig,
  ): SafeSuccess<T> | SafeFailure<T> {
    return run(argv, { ...config, safe: true }) as SafeSuccess<T> | SafeFailure<T>;
  }

  async function safeParseAsync(
    argv?: string[] | null,
    config?: ParseConfig,
  ): Promise<SafeSuccess<T> | SafeFailure<T>> {
    return run(argv, { ...config, safe: true }) as SafeSuccess<T> | SafeFailure<T>;
  }

  // function __getOptions(argv?: string[] | null) {
  //   const result = safeParse(argv);
  //   if (result.success) {
  //     return result.data.options as z.infer<typeof optionsSchema>;
  //   }

  //   return {} as z.infer<typeof optionsSchema>;
  // }
}

types.ts

import { z } from 'zod';

export type ExtractZodKeys<T> =
  T extends z.ZodObject<infer U> ? keyof z.infer<z.ZodObject<U>> : never;

export type FlagDefinition<T extends z.ZodObject<z.ZodRawShape>> = {
  [key in string]: ExtractZodKeys<T>;
} & {
  // This will ensure that keys from options schema cannot be used
  [key in ExtractZodKeys<T>]: never;
};

export type FlagsDefPartial<TT extends z.ZodObject<z.ZodRawShape>> = Partial<FlagDefinition<TT>>;

export type ParseConfig = { allowUnknown?: boolean; safe?: boolean };

export type DefineOptionsConfig<T extends z.ZodObject<z.ZodRawShape>> = {
  schema: T;
  aliases?: FlagsDefPartial<T>;
  raw?: boolean;
};
export type DefineCommandConfig<
  T extends z.ZodObject<z.ZodRawShape>,
  CommandArgs extends z.ZodTypeAny[] = z.ZodTypeAny[],
> = {
  name: string;
  description: string;
  args?: [...CommandArgs];
  aliases?: string[];
  options?: DefineOptionsConfig<T>;
  action: (
    // Exclude<ReturnType<ReturnType<typeof defineOptions>['parse']>, 'schema'>
    // Extract<ReturnType<ReturnType<typeof defineOptions>['parse']>, Pick<DefineOptionsConfig<T>, 'schema'>>
    commandOptions: z.infer<T>,
    args: { [K in keyof CommandArgs]: z.infer<CommandArgs[K]> },
  ) => void;
};

export type ResultData<T extends z.ZodObject<z.ZodRawShape>> = { schema: T; options: z.infer<T> };
export type SafeSuccess<T extends z.ZodObject<z.ZodRawShape>> = {
  success: true;
  data: ResultData<T>;
};
export type SafeFailure<T> = { success: false; error: z.ZodError & { schema?: T } };

export type StandardDefineReturn<T extends z.ZodObject<z.ZodRawShape>> = {
  run: (
    argv?: string[] | null,
    config?: ParseConfig,
  ) => ResultData<T> | SafeSuccess<T> | SafeFailure<T>;
  parse: (argv?: string[] | null, config?: ParseConfig) => ResultData<T>;
  parseAsync: (argv?: string[] | null, config?: ParseConfig) => Promise<ResultData<T>>;
  safeParse: (argv?: string[] | null, config?: ParseConfig) => SafeSuccess<T> | SafeFailure<T>;
  safeParseAsync: (
    argv?: string[] | null,
    config?: ParseConfig,
  ) => Promise<SafeSuccess<T> | SafeFailure<T>>;
};

utils.ts

import { z } from 'zod';

import type { DefineConfig, ParseConfig } from './types.ts';

export function parseLongFlag(arg: string): [string, string | null] {
  const parts = arg.slice(2).split('=', 2);
  let key = parts[0] || '';
  let value = parts[1] ?? null;

  // Handle --no-flag syntax
  if (key.startsWith('no-')) {
    key = key.slice(3);
    value = 'false';
  } else if (value === null) {
    value = 'true';
  }

  return [key, value];
}

export function parseShortFlags(arg: string): [string[], string | null] {
  const parts = arg.slice(1).split('=', 2);
  const flags = parts[0]?.split('') ?? [];
  const value = parts[1] ?? null;

  // For -abc=value, each flag gets the value
  // For -abc, each flag gets true
  return [flags, value ?? 'true'];
}

export function handleFlag<T extends z.ZodObject<z.ZodRawShape>>(
  { key, value }: { key: string; value: string | null },
  {
    schema,
    aliases,
    rawResults,
    unknownKeys,
  }: DefineConfig<T> & {
    rawResults: Record<string, string[]>;
    unknownKeys: string[];
  },
  { allowUnknown }: ParseConfig = { allowUnknown: false },
) {
  let flagName: string = key;

  // Check if it's a global flag
  if (aliases && key in aliases) {
    flagName = aliases[key] as string;
  }

  // Validate flag exists if not allowing unknown
  if (!allowUnknown && !(flagName in schema.shape)) {
    unknownKeys.push(key);
    return;
  }

  if (rawResults && flagName in rawResults) {
    rawResults[flagName]?.push(value ?? 'true');
  } else {
    rawResults[flagName] = [value ?? 'true'];
  }
}

export function parseFlags<T extends z.ZodObject<z.ZodRawShape>>(
  argv: string[],
  { schema, aliases }: DefineConfig<T>,
  { allowUnknown }: ParseConfig = { allowUnknown: false },
) {
  const rawResults: Record<string, string[]> = {};
  const unknownKeys: string[] = [];
  let cmdIndex = -1;

  for (let i = 0; i < argv.length; i++) {
    const arg = argv[i] || '';

    if (arg.startsWith('-')) {
      if (arg.startsWith('--')) {
        // Handle long flags (--flag)
        const [key, value] = parseLongFlag(arg);

        handleFlag({ key, value }, { schema, aliases, rawResults, unknownKeys }, { allowUnknown });
      } else {
        // Handle short flags (-a or -abc)
        const [keys, value] = parseShortFlags(arg);
        for (const key of keys) {
          handleFlag(
            { key, value },
            { schema, aliases, rawResults, unknownKeys },
            { allowUnknown },
          );
        }
      }
    } else {
      // Found command or argument
      cmdIndex = i;
      break;
    }
  }

  return { rawResults, cmdIndex, unknownKeys };
}

// Helper to transform raw results into final form
export function transformResults(results: Record<string, string[]>) {
  const transformed: Record<string, any> = {};

  for (const [key, values] of Object.entries(results)) {
    if (values.length === 1) {
      transformed[key] = values[0] === 'true' ? true : values[0] === 'false' ? false : values[0];
    } else {
      transformed[key] = values;
    }
  }

  return transformed;
}