From dcd5162710fcc7f354ffa5eaed574e0cd6f03d06 Mon Sep 17 00:00:00 2001
From: Paul Ouellette <oue.paul18@gmail.com>
Date: Mon, 27 May 2019 17:54:05 -0400
Subject: [PATCH] Add argument choices

---
 docsrc/arguments.rst    | 21 +++++++++++++++++++++
 docsrc/misc.rst         |  2 ++
 docsrc/options.rst      |  5 +++++
 spec/arguments_spec.lua | 10 ++++++++++
 spec/options_spec.lua   | 10 ++++++++++
 src/argparse.lua        | 16 ++++++++++++++++
 6 files changed, 64 insertions(+)

diff --git a/docsrc/arguments.rst b/docsrc/arguments.rst
index 02d65ef..9f41b31 100644
--- a/docsrc/arguments.rst
+++ b/docsrc/arguments.rst
@@ -68,3 +68,24 @@ If more than one argument can be consumed, a table is used to store the data.
       pair = {"foo", "bar"},
       optional = "baz"
    }
+
+Setting argument choices
+------------------------
+
+The ``choices`` property can be used to restrict an argument to a set of choices. Its value is an array of string choices.
+
+.. code-block:: lua
+   :linenos:
+
+   parser:argument "direction"
+      :choices {"north", "south", "east", "west"}
+
+.. code-block:: none
+
+   $ lua script.lua foo
+
+.. code-block:: none
+
+   Usage: script.lua [-h] <direction>
+
+   Error: argument 'direction' must be one of 'north', 'south', 'east', 'west'
diff --git a/docsrc/misc.rst b/docsrc/misc.rst
index bdd74e4..05a42ab 100644
--- a/docsrc/misc.rst
+++ b/docsrc/misc.rst
@@ -242,6 +242,7 @@ Property            Type
 ``defmode``         String
 ``show_default``    Boolean
 ``argname``         String or table
+``choices``         Table
 ``action``          Function or string
 ``init``            Any
 ``hidden``          Boolean
@@ -273,6 +274,7 @@ Property            Type
 ``show_default``    Boolean
 ``overwrite``       Booleans
 ``argname``         String or table
+``choices``         Table
 ``action``          Function or string
 ``init``            Any
 ``hidden``          Boolean
diff --git a/docsrc/options.rst b/docsrc/options.rst
index bcc7556..f88bc04 100644
--- a/docsrc/options.rst
+++ b/docsrc/options.rst
@@ -107,6 +107,11 @@ Just as arguments, options can be configured to take several command line argume
 
 Note that the data passed to ``optional`` option is stored in an array. That is necessary to distinguish whether the option was invoked without an argument or it was not invoked at all.
 
+Setting argument choices
+------------------------
+
+The ``choices`` property can be used to specify a list of choices for an option argument in the same way as for arguments.
+
 Setting number of invocations
 -----------------------------
 
diff --git a/spec/arguments_spec.lua b/spec/arguments_spec.lua
index 8a1dbef..5cf09c4 100644
--- a/spec/arguments_spec.lua
+++ b/spec/arguments_spec.lua
@@ -161,5 +161,15 @@ describe("tests related to positional arguments", function()
          }
          assert.has_error(function() parser:parse{} end, "missing argument 'foo1'")
       end)
+
+      it("handles invalid argument choices correctly", function()
+         local parse = Parser()
+         parse:argument "foo" {
+            choices = {"bar", "baz", "qu"}
+         }
+         assert.has_error(function()
+            parse:parse{"foo", "quu"}
+         end, "argument 'foo' must be one of 'bar', 'baz', 'qu'")
+      end)
    end)
 end)
diff --git a/spec/options_spec.lua b/spec/options_spec.lua
index 40fa4b0..4fdffe8 100644
--- a/spec/options_spec.lua
+++ b/spec/options_spec.lua
@@ -290,6 +290,16 @@ describe("tests related to options", function()
          end, "too many arguments")
       end)
 
+      it("handles invalid argument choices correctly", function()
+         local parse = Parser()
+         parse:option "-s" "--server" {
+            choices = {"foo", "bar", "baz"}
+         }
+         assert.has_error(function()
+            parse:parse{"-slocalhost"}
+         end, "argument for option '-s' must be one of 'foo', 'bar', 'baz'")
+      end)
+
       it("doesn't accept GNU-like long options when it doesn't need arguments", function()
          local parser = Parser()
          parser:flag "-q" "--quiet"
diff --git a/src/argparse.lua b/src/argparse.lua
index 91ab584..443478d 100644
--- a/src/argparse.lua
+++ b/src/argparse.lua
@@ -341,6 +341,7 @@ local Argument = class({
    typechecked("defmode", "string"),
    typechecked("show_default", "boolean"),
    typechecked("argname", "string", "table"),
+   typechecked("choices", "table"),
    typechecked("hidden", "boolean"),
    option_action,
    option_init
@@ -363,6 +364,7 @@ local Option = class({
    typechecked("show_default", "boolean"),
    typechecked("overwrite", "boolean"),
    typechecked("argname", "string", "table"),
+   typechecked("choices", "table"),
    typechecked("hidden", "boolean"),
    option_action,
    option_init
@@ -1204,7 +1206,21 @@ function ElementState:invoke()
    return self.open
 end
 
+function ElementState:check_choices(argument)
+   if self.element._choices then
+      for _, choice in ipairs(self.element._choices) do
+         if argument == choice then
+            return
+         end
+      end
+      local choices_list = "'" .. table.concat(self.element._choices, "', '") .. "'"
+      local is_option = getmetatable(self.element) == Option
+      self:error("%s%s must be one of %s", is_option and "argument for " or "", self.name, choices_list)
+   end
+end
+
 function ElementState:pass(argument)
+   self:check_choices(argument)
    argument = self:convert(argument, #self.args + 1)
    table.insert(self.args, argument)