diff --git a/argparse_prompt.py b/argparse_prompt.py index bf98395..15c3d19 100644 --- a/argparse_prompt.py +++ b/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