feat(core): support enum for argv type, fix #337

koishijs · Aug 15, 2021 · 3b2adf9 · 3b2adf9
1 parent 792233a
commit 3b2adf9
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 12 deletions.
diff --git a/packages/koishi-core/src/command.ts b/packages/koishi-core/src/command.ts
@@ -173,10 +173,15 @@ export class Command<U extends User.Field = never, G extends Channel.Field = nev
     return this
   }
 
-  option<K extends string, D extends string, T extends Argv.Type>(name: K, desc: D, config: Argv.OptionConfig<T> = {}) {
+  option<K extends string, T extends keyof Argv.Domain>(name: K, desc: string, config: Argv.OptionConfig<T>): Command<U, G, A, Extend<O, K, Argv.Domain[T]>>
+  option<K extends string, R extends string>(name: K, desc: string, config: Argv.OptionConfig<R[]>): Command<U, G, A, Extend<O, K, R>>
+  option<K extends string, R>(name: K, desc: string, config: Argv.OptionConfig<(source: string) => R>): Command<U, G, A, Extend<O, K, R>>
+  option<K extends string>(name: K, desc: string, config: Argv.OptionConfig<RegExp>): Command<U, G, A, Extend<O, K, string>>
+  option<K extends string, D extends string>(name: K, desc: D, config?: Argv.OptionConfig): Command<U, G, A, Extend<O, K, Argv.OptionType<D>>>
+  option(name: string, desc: string, config: Argv.OptionConfig = {}) {
     this._createOption(name, desc, config)
     this._disposables?.push(() => this.removeOption(name))
-    return this as Command<U, G, A, Extend<O, K, Argv.OptionType<D, T>>>
+    return this
   }
 
   match(session: Session) {

diff --git a/packages/koishi-core/src/parser.ts b/packages/koishi-core/src/parser.ts
@@ -187,15 +187,9 @@ export namespace Argv {
 
   export type ArgumentType<S extends string> = ExtractSpread<Replace<S, '>', ']'>>
 
-  // I don't know why I should write like this but
-  // [T] extends [xxx] just works, so don't touch it
-  export type OptionType<S extends string, T extends Type>
-    = [T] extends [DomainType] ? Domain[T]
-    : [T] extends [RegExp] ? string
-    : T extends (source: string) => infer R ? R
-    : ExtractFirst<Replace<S, '>', ']'>, any>
+  export type OptionType<S extends string> = ExtractFirst<Replace<S, '>', ']'>, any>
 
-  export type Type = DomainType | RegExp | Transform<any>
+  export type Type = DomainType | RegExp | string[] | Transform<any>
 
   export interface Declaration {
     name?: string
@@ -217,12 +211,18 @@ export namespace Argv {
   }
 
   function resolveType(type: Type) {
-    if (typeof type === 'function') return type
-    if (type instanceof RegExp) {
+    if (typeof type === 'function') {
+      return type
+    } else if (type instanceof RegExp) {
       return (source: string) => {
         if (type.test(source)) return source
         throw new Error()
       }
+    } else if (Array.isArray(type)) {
+      return (source: string) => {
+        if (type.includes(source)) return source
+        throw new Error()
+      }
     }
     return builtin[type]?.transform
   }

diff --git a/packages/koishi-core/tests/runtime.spec.ts b/packages/koishi-core/tests/runtime.spec.ts
@@ -265,10 +265,12 @@ describe('Runtime', () => {
       cmd3.option('foo', '<foo>', { type: () => { throw new Error() } })
       cmd3.option('bar', '<bar>', { type: () => { throw new Error('SUFFIX') } })
       cmd3.option('baz', '<baz>', { type: /$^/ })
+      cmd3.option('bax', '<baz>', { type: ['abc', 'def'] })
       await session1.shouldReply('cmd3', 'after cmd3')
       await session1.shouldReply('cmd3 --foo xxx', '选项 foo 输入无效，输入帮助以查看用法。')
       await session1.shouldReply('cmd3 --bar xxx', '选项 bar 输入无效，SUFFIX')
       await session1.shouldReply('cmd3 --baz xxx', '选项 baz 输入无效，输入帮助以查看用法。')
+      await session1.shouldReply('cmd3 --bax cba', '选项 bax 输入无效，输入帮助以查看用法。')
       cmd3.dispose()
     })