Added support for parameter choices

argparse_prompt.py

@@ -9,7 +9,7 @@ class PromptParser(argparse.ArgumentParser):
     Extends ArgumentParser to allow any unspecified arguments to be input dynamically on the command line
     """
 
-    def add_argument(self, *args, prompt=True, secure=False, **kwargs):
+    def add_argument(self, *args, prompt=True, choices=False, secure=False, **kwargs):
         """
         For all unlisted arguments, refer to the parent class
         :param prompt: False if we never want to prompt the user for this argument
@@ -27,6 +27,7 @@ def add_argument(self, *args, prompt=True, secure=False, **kwargs):
                 type=kwargs.get("type"),
                 secure=secure,
                 default=kwargs.get("default"),
+                choices=choices
             )
 
             # Remove the old type so we can replace it with our own
@@ -36,7 +37,10 @@ def add_argument(self, *args, prompt=True, secure=False, **kwargs):
                 del kwargs["default"]
 
             # Delegate to the parent class. Default must be '' in order to get the type function to be called
-            action = super().add_argument(*args, type=type, default="", **kwargs)
+            if(choices):
+                action = super().add_argument(*args, type=type, choices=choices, default="", **kwargs)
+            else:
+                action = super().add_argument(*args, type=type, default="", **kwargs)
 
             # Set the argument name, now that the parser has parsed it
             type.name = action.dest
@@ -50,7 +54,7 @@ class Prompt:
     A class the pretends to be a function so that it can be used as the 'type' argument for the ArgumentParser
     """
 
-    def __init__(self, name=None, help=None, type=None, default=None, secure=False):
+    def __init__(self, name=None, help=None, type=None, default=None, secure=False, choices=False):
         """
         Creates a new prompt validator
         :param name: The identifier for the variable
@@ -64,20 +68,36 @@ def __init__(self, name=None, help=None, type=None, default=None, secure=False):
         self.help = help
         self.default = default
         self.secure = secure
+        self.choices = choices
 
     def __call__(self, val):
         default_str = f"({self.default}) "
         help_str = "" if self.help is None else f": {self.help}"
+
+        # Make our choices pretty for display
+        choices = self.choices
+        if type(choices) == list:
+            choices_options ="({options}): ".format(options='|'.join(choices))
+        else:
+            choices_options = False
 
         try:
             # If the user provided no value for this argument, prompt them for it
             if val == "":
-                prompt = "{}{}\n> {}".format(self.name, help_str, default_str)
+                prompt = "{}{}\n> {}".format(self.name, help_str, choices_options)
 
                 newval = (
                     getpass.getpass(prompt=prompt) if self.secure else input(prompt)
                 )
 
+                while choices_options and newval not in choices:
+                    help_str = "You must choose a {} from the provided options shown below:".format(self.name)
+                    prompt = "{}\n> {}".format( help_str, choices_options)
+
+                    newval = (
+                        getpass.getpass(prompt=prompt) if self.secure else input(prompt)
+                    )
+
                 # If they just hit enter, they want the default value
                 if newval == "":
                     newval = self.default