giuspen giuspen - 1 year ago 240
Python Question

python argparse choices of string list accept unique partial list element

I would like to have the following rule

parser.add_argument('move', choices=['rock', 'paper', 'scissors'])

work also if you pass an unique subset of the characters (e.g. "k" or "oc" will be accepted as "rock" but not "r" since it's not unique).

My need is to be able to run the script with one or more arguments in the quickest possible way so avoiding writing the whole argument name when a subset would be sufficient for the script to understand the choice.

Is there a way to obtain this result still exploiting the handy choices list which integrates automatically in the help and in the errors handling?

Answer Source

You could define a custom subclass of list that supports your definition of the in operator as "contains (part of) an element exactly once" like this:

class argList(list):
    def __contains__(self, other):
        "Check if <other> is a substring of exactly one element of <self>"
        num = 0
        for item in self:
            if other in item:
                num += 1
            if num > 1:
                return False
        return num==1
    def find(self, other):
        "Return the first element of <self> in which <other> can be found"
        for item in self:
            if other in item:
                return item

You could then use it to construct your argument list:

>>> l = argList(["rock", "paper", "scissors"])
>>> "rock" in l
>>> "ck" in l
>>> "r" in l
>>> "q" in l

and use it to construct your parser:

>>> import argparse
>>> parser = argparse.ArgumentParser()
>>> l = argList(["rock", "paper", "scissors"])
>>> parser.add_argument("move", choices=l)
_StoreAction(option_strings=[], dest='move', nargs=None, const=None, default=Non
e, type=None, choices=['rock', 'paper', 'scissors'], help=None, metavar=None)

Now it handles arguments correctly (although the error message is still a bit misleading):

>>> parser.parse_args(["rock"])
>>> parser.parse_args(["r"])
usage: [-h] {rock,paper,scissors}
: error: argument move: invalid choice: 'r' (choose from 'rock', 'paper', 'scissors')

Note that the entered argument is retained which (of course) may be incomplete:

>>> parser.parse_args(["ck"])

so you will have to find out which of your actual arguments was chosen:

>>> args = parser.parse_args(["ck"])
>>> l.find(vars(args)["move"])

I haven't tested this much further - it might be that redefining in` has unexpected side effects in argument lists, and I guess it's questionable whether such program behavior violates the principle of least surprise - but it's a start.