From 091a17b10ba91df87338b0144be1ed6cbeaaaae8 Mon Sep 17 00:00:00 2001 From: levontumanyan <5532305+levontumanyan@users.noreply.github.com> Date: Wed, 17 Jun 2026 19:33:14 +0200 Subject: [PATCH] feat(completion): complete enum flag values from schema Extends the completion engine to offer valid enum values as tab candidates when the user is typing a value for an enum-typed flag (e.g. elastic stack es cat nodes --h ). How it works: - extractSchemaArgs() now extracts enumValues from z.enum (and union(z.enum, z.string)) fields via a new extractEnumValues() helper. - defineCommand() attaches a non-enumerable _enumCompleters map to the Commander command for each flag that has enumValues. - enumerate() falls back to _enumCompleters when the global registry has no completer for the previous flag token. Both the --flag value and --flag=partial forms are handled. Since enum values come from the generated Zod schemas, all cat API column flags (--h, --s where typed as enum) and any other enum-typed flags across the codebase get completions for free without any manual registration. --- src/completion/enumerate.ts | 6 +++-- src/factory.ts | 11 ++++++++ src/lib/schema-args.ts | 38 +++++++++++++++++++++++++- test/completion/enumerate.test.ts | 44 +++++++++++++++++++++++++++++++ test/lib/schema-args.test.ts | 34 ++++++++++++++++++++++++ 5 files changed, 130 insertions(+), 3 deletions(-) diff --git a/src/completion/enumerate.ts b/src/completion/enumerate.ts index ae26cc56..9c5685a4 100644 --- a/src/completion/enumerate.ts +++ b/src/completion/enumerate.ts @@ -186,7 +186,8 @@ export async function enumerate ( const eqIdx = incomplete.indexOf('=') const flag = incomplete.slice(0, eqIdx) const partial = incomplete.slice(eqIdx + 1) - const completer = registry?.get(flag) + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const completer = registry?.get(flag) ?? (current as any)._enumCompleters?.get(flag) as DynamicCompleter | undefined if (completer != null) { const cands = await safeRun(completer) return { @@ -206,7 +207,8 @@ export async function enumerate ( } if (previous != null && previous.startsWith('--')) { - const completer = registry?.get(previous) + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const completer = registry?.get(previous) ?? (current as any)._enumCompleters?.get(previous) as DynamicCompleter | undefined if (completer != null) { const cands = await safeRun(completer) return { diff --git a/src/factory.ts b/src/factory.ts index 04ce03b3..b1abf658 100644 --- a/src/factory.ts +++ b/src/factory.ts @@ -366,6 +366,17 @@ export function defineCommand (config: CommandConfig): O } else if (arg.type === 'enum') { const attrName = camelCase(arg.cliFlag) cmd.option(`--${arg.cliFlag} `, desc, singleValueGuard(cmd, attrName, `--${arg.cliFlag}`)) + if (arg.enumValues != null && arg.enumValues.length > 0) { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + let completers = (cmd as any)._enumCompleters as Map string[]> | undefined + if (completers == null) { + completers = new Map() + // eslint-disable-next-line @typescript-eslint/no-explicit-any + Object.defineProperty(cmd as any, '_enumCompleters', { value: completers, enumerable: false }) + } + const vals = arg.enumValues + completers.set(`--${arg.cliFlag}`, () => Array.from(vals)) + } } else { // string: accumulate repeated values with comma separation const attrName = camelCase(arg.cliFlag) diff --git a/src/lib/schema-args.ts b/src/lib/schema-args.ts index fcac9590..5a9b7a2f 100644 --- a/src/lib/schema-args.ts +++ b/src/lib/schema-args.ts @@ -37,6 +37,12 @@ export interface SchemaArgDefinition { */ acceptsArrayForm?: boolean + /** + * Valid enum values extracted from the Zod schema, present when `type === 'enum'`. + * Used by the completion engine to offer candidates for this flag without a live cluster. + */ + enumValues?: readonly string[] + /** * Marks args whose CLI string value needs a non-trivial transformation before reaching the wire. * @@ -87,6 +93,8 @@ interface ZodFieldDef { defaultValue?: unknown getter?: () => z.ZodType options?: z.ZodType[] + /** Present on `z.enum(...)` nodes — maps each value to itself. */ + entries?: Record } /** @@ -124,6 +132,32 @@ function unwrapField (field: z.ZodType): { typeName: string, isOptional: boolean return { typeName: def.type, isOptional: false } } +/** + * Walks a Zod field through wrapper types (`optional`, `default`, `lazy`, `union`) and + * returns the enum values if any branch resolves to a `z.enum(...)` node. + * + * Returns `undefined` for non-enum fields. Never throws. + */ +function extractEnumValues (field: z.ZodType): readonly string[] | undefined { + const def = field.def as ZodFieldDef + if (def.type === 'enum' && def.entries != null) { + return Object.values(def.entries) + } + if ((def.type === 'optional' || def.type === 'default') && def.innerType != null) { + return extractEnumValues(def.innerType as z.ZodType) + } + if (def.type === 'lazy' && typeof def.getter === 'function') { + return extractEnumValues(def.getter()) + } + if (def.type === 'union' && Array.isArray(def.options)) { + for (const opt of def.options) { + const vals = extractEnumValues(opt) + if (vals != null) return vals + } + } + return undefined +} + const CLI_TYPES = new Set(['string', 'number', 'boolean', 'object', 'array', 'enum']) /** @@ -158,6 +192,7 @@ export function extractSchemaArgs (schema: unknown): SchemaArgDefinition[] { const foundIn = extractFoundIn(fieldSchema as z.ZodType) const acceptsArrayForm = type !== 'array' && schemaAcceptsArrayForm(fieldSchema as z.ZodType) const parseStyle = schemaContainsId(fieldSchema as z.ZodType, 'Sort') ? 'sort-pairs' as const : undefined + const enumValues = type === 'enum' ? extractEnumValues(fieldSchema as z.ZodType) : undefined return { schemaKey: key, cliFlag: toKebabCase(key), @@ -167,7 +202,8 @@ export function extractSchemaArgs (schema: unknown): SchemaArgDefinition[] { description, ...(foundIn !== undefined ? { foundIn } : {}), ...(acceptsArrayForm ? { acceptsArrayForm: true } : {}), - ...(parseStyle !== undefined ? { parseStyle } : {}) + ...(parseStyle !== undefined ? { parseStyle } : {}), + ...(enumValues !== undefined ? { enumValues } : {}) } }) } diff --git a/test/completion/enumerate.test.ts b/test/completion/enumerate.test.ts index 976e28ba..84ef736d 100644 --- a/test/completion/enumerate.test.ts +++ b/test/completion/enumerate.test.ts @@ -261,3 +261,47 @@ describe('enumerate -- edge cases', () => { assert.deepEqual(result.candidates, ['stack']) }) }) + +describe('enumerate -- command-level enum completers (_enumCompleters)', () => { + function buildProgramWithEnumFlag (): Command { + const program = buildSampleProgram() + const catNodes = defineCommand({ + name: 'nodes', + description: 'Cat nodes', + handler: () => ({}), + }) + // Simulate what factory attaches for enum flags + const completers = new Map string[]>([ + ['--h', () => ['cpu', 'ram', 'name']], + ]) + Object.defineProperty(catNodes, '_enumCompleters', { value: completers, enumerable: false }) + const cat = defineGroup({ name: 'cat', description: 'Cat APIs' }, catNodes) + const es = defineGroup({ name: 'es', description: 'ES' }, cat) + program.addCommand(es) + return program + } + + it('returns enum values when previous word is a flag with a command-level completer', async () => { + const result = await enumerate(buildProgramWithEnumFlag(), ['es', 'cat', 'nodes', '--h', '']) + assert.deepEqual(result.candidates, ['cpu', 'ram', 'name']) + }) + + it('prefix-filters enum candidates', async () => { + const result = await enumerate(buildProgramWithEnumFlag(), ['es', 'cat', 'nodes', '--h', 'c']) + assert.deepEqual(result.candidates, ['cpu']) + }) + + it('supports --flag=partial form for command-level enum completers', async () => { + const result = await enumerate(buildProgramWithEnumFlag(), ['es', 'cat', 'nodes', '--h=ra']) + assert.deepEqual(result.candidates, ['ram']) + assert.equal(result.directive & DIRECTIVE_NO_SPACE, DIRECTIVE_NO_SPACE) + }) + + it('global registry completer takes priority over command-level enum completer', async () => { + const registry: DynamicCompleterRegistry = { + get: (flag) => flag === '--h' ? () => ['override'] : undefined, + } + const result = await enumerate(buildProgramWithEnumFlag(), ['es', 'cat', 'nodes', '--h', ''], registry) + assert.deepEqual(result.candidates, ['override']) + }) +}) diff --git a/test/lib/schema-args.test.ts b/test/lib/schema-args.test.ts index f3a35ee2..a7f6f761 100644 --- a/test/lib/schema-args.test.ts +++ b/test/lib/schema-args.test.ts @@ -316,3 +316,37 @@ describe('extractSchemaArgs acceptsArrayForm detection (#167)', () => { assert.notEqual(args[0]?.acceptsArrayForm, true) }) }) + +describe('extractSchemaArgs -- enumValues', () => { + it('populates enumValues for a plain z.enum field', () => { + const schema = z.object({ col: z.enum(['cpu', 'ram', 'name']) }) + const args = extractSchemaArgs(schema) + assert.deepEqual(args[0]?.enumValues, ['cpu', 'ram', 'name']) + }) + + it('populates enumValues through optional wrapper', () => { + const schema = z.object({ col: z.enum(['cpu', 'ram']).optional() }) + const args = extractSchemaArgs(schema) + assert.deepEqual(args[0]?.enumValues, ['cpu', 'ram']) + }) + + it('populates enumValues from union(enum, string) — ES cat column pattern', () => { + const CatColumn = z.union([z.enum(['cpu', 'ram', 'name']), z.string()]) + const schema = z.object({ h: CatColumn.optional() }) + const args = extractSchemaArgs(schema) + assert.deepEqual(args[0]?.enumValues, ['cpu', 'ram', 'name']) + assert.equal(args[0]?.type, 'enum') + }) + + it('leaves enumValues undefined for plain string fields', () => { + const schema = z.object({ name: z.string() }) + const args = extractSchemaArgs(schema) + assert.equal(args[0]?.enumValues, undefined) + }) + + it('leaves enumValues undefined for number fields', () => { + const schema = z.object({ count: z.number() }) + const args = extractSchemaArgs(schema) + assert.equal(args[0]?.enumValues, undefined) + }) +})