feat: support enum arg type (#83)

unjs · Mar 18, 2024 · 02f4fdb · 02f4fdb
1 parent 34c52a2
commit 02f4fdb
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 6 deletions.
diff --git a/playground/hello.ts b/playground/hello.ts
@@ -17,9 +17,16 @@ const command = defineCommand({
       type: "boolean",
       description: "Use friendly greeting",
     },
+    adjective: {
+      type: "enum",
+      description: "Adjective to use in greeting",
+      options: ["awesome", "cool", "nice"],
+      default: "awesome",
+      require: false,
+    }
   },
   run({ args }) {
-    consola.log(`${args.friendly ? "Hi" : "Greetings"} ${args.name}!`);
+    consola.log(`${args.friendly ? "Hi" : "Greetings"} ${args.adjective} ${args.name}!`);
   },
 });
 

diff --git a/src/args.ts b/src/args.ts
@@ -10,6 +10,7 @@ export function parseArgs<T extends ArgsDef = ArgsDef>(
   const parseOptions = {
     boolean: [] as string[],
     string: [] as string[],
+    enum: [] as (number | string)[],
     mixed: [] as string[],
     alias: {} as Record<string, string | string[]>,
     default: {} as Record<string, boolean | string>,
@@ -21,11 +22,15 @@ export function parseArgs<T extends ArgsDef = ArgsDef>(
     if (arg.type === "positional") {
       continue;
     }
+    // eslint-disable-next-line unicorn/prefer-switch
     if (arg.type === "string") {
       parseOptions.string.push(arg.name);
     } else if (arg.type === "boolean") {
       parseOptions.boolean.push(arg.name);
+    } else if (arg.type === "enum") {
+      parseOptions.enum.push(...(arg.options || []));
     }
+
     if (arg.default !== undefined) {
       parseOptions.default[arg.name] = arg.default;
     }
@@ -56,6 +61,19 @@ export function parseArgs<T extends ArgsDef = ArgsDef>(
       } else {
         parsedArgsProxy[arg.name] = arg.default;
       }
+    } else if (arg.type === "enum") {
+      const argument = parsedArgsProxy[arg.name];
+      const options = arg.options || [];
+      if (
+        argument !== undefined &&
+        options.length > 0 &&
+        !options.includes(argument)
+      ) {
+        throw new CLIError(
+          `Invalid value for argument: \`--${arg.name}\` (\`${argument}\`). Expected one of: ${options.map((o) => `\`${o}\``).join(", ")}.`,
+          "EARG",
+        );
+      }
     } else if (arg.required && parsedArgsProxy[arg.name] === undefined) {
       throw new CLIError(`Missing required argument: --${arg.name}`, "EARG");
     }

diff --git a/src/types.ts b/src/types.ts
@@ -1,6 +1,6 @@
 // ----- Args -----
 
-export type ArgType = "boolean" | "string" | "positional" | undefined;
+export type ArgType = "boolean" | "string" | "positional" | "enum" | undefined;
 
 export type _ArgDef<T extends ArgType, VT extends boolean | string> = {
   type?: T;
@@ -9,12 +9,21 @@ export type _ArgDef<T extends ArgType, VT extends boolean | string> = {
   alias?: string | string[];
   default?: VT;
   required?: boolean;
+  options?: (string | number)[];
 };
 
-export type BooleanArgDef = _ArgDef<"boolean", boolean>;
-export type StringArgDef = _ArgDef<"string", string>;
-export type PositionalArgDef = Omit<_ArgDef<"positional", string>, "alias">;
-export type ArgDef = BooleanArgDef | StringArgDef | PositionalArgDef;
+export type BooleanArgDef = Omit<_ArgDef<"boolean", boolean>, "options">;
+export type StringArgDef = Omit<_ArgDef<"string", string>, "options">;
+export type PositionalArgDef = Omit<
+  _ArgDef<"positional", string>,
+  "alias" | "options"
+>;
+export type EnumArgDef = _ArgDef<"enum", string>;
+export type ArgDef =
+  | BooleanArgDef
+  | StringArgDef
+  | PositionalArgDef
+  | EnumArgDef;
 export type ArgsDef = Record<string, ArgDef>;
 export type Arg = ArgDef & { name: string; alias: string[] };
 
@@ -34,6 +43,14 @@ export type ParsedArgs<T extends ArgsDef = ArgsDef> = { _: string[] } & Record<
     }[keyof T],
     boolean
   > &
+  Record<
+    {
+      [K in keyof T]: T[K] extends { type: "enum" } ? K : never;
+    }[keyof T],
+    {
+      [K in keyof T]: T[K] extends { options: infer U } ? U : never;
+    }[keyof T]
+  > &
   Record<string, string | boolean | string[]>;
 
 // ----- Command -----

diff --git a/src/usage.ts b/src/usage.ts
@@ -59,6 +59,9 @@ export async function renderUsage<T extends ArgsDef = ArgsDef>(
           ? `=${
               arg.valueHint ? `<${arg.valueHint}>` : `"${arg.default || ""}"`
             }`
+          : "") +
+        (arg.type === "enum" && arg.options
+          ? `=<${arg.options.join("|")}>`
           : "");
       argLines.push([
         "`" + argStr + (isRequired ? " (required)" : "") + "`",