From 6a33b0633c9e422b18ec33d9a1dd4202c97adf28 Mon Sep 17 00:00:00 2001
From: Jeff Dickey <216188+jdxcode@users.noreply.github.com>
Date: Wed, 31 Jan 2018 20:29:07 -0800
Subject: [PATCH] fix: added options to flags/args

---
 src/args.ts        |  2 ++
 src/errors.ts      | 28 +++++++++++++++++++++++-----
 src/flags.ts       | 12 +-----------
 src/parse.ts       | 12 ++++++++++--
 test/parse.test.ts | 29 ++++++++++++++++++++++-------
 5 files changed, 58 insertions(+), 25 deletions(-)

diff --git a/src/args.ts b/src/args.ts
index 2ae5c59..40d56e8 100644
--- a/src/args.ts
+++ b/src/args.ts
@@ -7,6 +7,7 @@ export interface IArg<T = string> {
   hidden?: boolean
   parse?: ParseFn<T>
   default?: T | (() => T)
+  options?: string[]
 }
 
 export interface ArgBase<T> {
@@ -16,6 +17,7 @@ export interface ArgBase<T> {
   parse: ParseFn<T>
   default?: T | (() => T)
   input?: string
+  options?: string[]
 }
 
 export type RequiredArg<T> = ArgBase<T> & {
diff --git a/src/errors.ts b/src/errors.ts
index d10ae9f..f0a2709 100644
--- a/src/errors.ts
+++ b/src/errors.ts
@@ -1,13 +1,13 @@
 import {Arg} from './args'
 import {deps} from './deps'
-import {IFlag} from './flags'
+import * as flags from './flags'
 import {flagUsages} from './help'
 import {ParserInput, ParserOutput} from './parse'
 
 export interface ICLIParseErrorOptions {
   parse: {
-    input: ParserInput
-    output: ParserOutput<any, any>
+    input?: ParserInput
+    output?: ParserOutput<any, any>
   }
 }
 
@@ -37,9 +37,9 @@ export class RequiredArgsError extends CLIParseError {
 }
 
 export class RequiredFlagError extends CLIParseError {
-  public flags: IFlag<any>[]
+  public flags: flags.IFlag<any>[]
 
-  constructor({flags, parse}: ICLIParseErrorOptions & { flags: IFlag<any>[] }) {
+  constructor({flags, parse}: ICLIParseErrorOptions & { flags: flags.IFlag<any>[] }) {
     const usage = deps.renderList(flagUsages(flags, {displayRequired: false}))
     const message = `Missing required flag${flags.length === 1 ? '' : 's'}:\n${usage}`
     super({parse, message})
@@ -56,3 +56,21 @@ export class UnexpectedArgsError extends CLIParseError {
     this.args = args
   }
 }
+
+export class FlagInvalidOptionError extends CLIParseError {
+  public args: string[]
+
+  constructor(flag: flags.IOptionFlag<any>, input: string) {
+    const message = `Expected --${flag.name}=${input} to be one of: ${flag.options!.join(', ')}`
+    super({parse: {}, message})
+  }
+}
+
+export class ArgInvalidOptionError extends CLIParseError {
+  public args: string[]
+
+  constructor(arg: Arg<any>, input: string) {
+    const message = `Expected ${input} to be one of: ${arg.options!.join(', ')}`
+    super({parse: {}, message})
+  }
+}
diff --git a/src/flags.ts b/src/flags.ts
index 4777634..c54ca01 100644
--- a/src/flags.ts
+++ b/src/flags.ts
@@ -25,6 +25,7 @@ export type IOptionFlag<T> = IFlagBase<T, string> & {
   default?: T | ((context: DefaultContext<T>) => T | undefined)
   multiple: boolean
   input: string[]
+  options?: string[]
 }
 
 export type Definition<T> = {
@@ -71,17 +72,6 @@ export const integer = build({
   },
 })
 
-const _enum = <T = string>(opts: EnumFlagOptions<T>) => build<T>({
-  parse(input) {
-    if (!opts.options.includes(input)) throw new Error(`Expected --${this.name}=${input} to be one of: ${opts.options.join(', ')}`)
-    return input
-  },
-  helpValue: `(${opts.options.join('|')})`,
-  ...opts as any,
-  optionType: 'enum',
-})
-export {_enum as enum}
-
 export function option<T>(options: {parse: IOptionFlag<T>['parse']} & Partial<IOptionFlag<T>>) {
   return build<T>({optionType: 'custom', ...options})()
 }
diff --git a/src/parse.ts b/src/parse.ts
index b21571c..049eafd 100644
--- a/src/parse.ts
+++ b/src/parse.ts
@@ -3,6 +3,7 @@
 import * as _ from 'lodash'
 
 import {Arg} from './args'
+import * as Errors from './errors'
 import * as Flags from './flags'
 
 let debug: any
@@ -41,7 +42,7 @@ export class Parser<T extends ParserInput, TFlags extends OutputFlags<T['flags']
   private readonly argv: string[]
   private readonly raw: ParsingToken[] = []
   private readonly booleanFlags: { [k: string]: Flags.IBooleanFlag<any> }
-  constructor(readonly input: T) {
+  constructor(private readonly input: T) {
     this.argv = input.argv.slice(0)
     this._setNames()
     this.booleanFlags = _.pickBy(input.flags, f => f.type === 'boolean') as any
@@ -155,7 +156,11 @@ export class Parser<T extends ParserInput, TFlags extends OutputFlags<T['flags']
           flags[token.flag] = true
         }
       } else {
-        const value = flag.parse ? flag.parse(token.input) : token.input
+        const input = token.input
+        if (flag.options && !flag.options.includes(input)) {
+          throw new Errors.FlagInvalidOptionError(flag, input)
+        }
+        const value = flag.parse ? flag.parse(input) : input
         if (flag.multiple) {
           flags[token.flag] = flags[token.flag] || []
           flags[token.flag].push(value)
@@ -185,6 +190,9 @@ export class Parser<T extends ParserInput, TFlags extends OutputFlags<T['flags']
       const arg = this.input.args[i]
       if (token) {
         if (arg) {
+          if (arg.options && !arg.options.includes(token.input)) {
+            throw new Errors.ArgInvalidOptionError(arg, token.input)
+          }
           args[i] = arg.parse(token.input)
         } else {
           args[i] = token.input
diff --git a/test/parse.test.ts b/test/parse.test.ts
index fe597a2..b806b9f 100644
--- a/test/parse.test.ts
+++ b/test/parse.test.ts
@@ -338,22 +338,37 @@ See more help with --help`)
     })
   })
 
-  describe('enum', () => {
+  describe('flag options', () => {
     it('accepts valid option', () => {
-      const foo = flags.enum<'myopt' | 'myotheropt'>({options: ['myopt', 'myotheropt']})
       const out = parse(['--foo', 'myotheropt'], {
-        flags: {foo: foo()},
+        flags: {foo: flags.string({options: ['myopt', 'myotheropt']})},
       })
       expect(out.flags.foo).to.equal('myotheropt')
     })
 
     it('fails when invalid', () => {
-      const foo = flags.enum({options: ['myopt', 'myotheropt']})
       expect(() => {
-        parse(['--foo', 'bar'], {
-          flags: {foo: foo()},
+        parse(['--foo', 'invalidopt'], {
+          flags: {foo: flags.string({options: ['myopt', 'myotheropt']})},
+        })
+      }).to.throw('Expected --foo=invalidopt to be one of: myopt, myotheropt')
+    })
+  })
+
+  describe('arg options', () => {
+    it('accepts valid option', () => {
+      const out = parse(['myotheropt'], {
+        args: [{name: 'foo', options: ['myopt', 'myotheropt']}],
+      })
+      expect(out.args.foo).to.equal('myotheropt')
+    })
+
+    it('fails when invalid', () => {
+      expect(() => {
+        parse(['invalidopt'], {
+          args: [{name: 'foo', options: ['myopt', 'myotheropt']}],
         })
-      }).to.throw('Expected --foo=bar to be one of: myopt, myotheropt')
+      }).to.throw('Expected invalidopt to be one of: myopt, myotheropt')
     })
   })
 })