Python argparse choices from an infinite set

问题:

I have the following code to create a container which pretends to behave like the set of all prime numbers (actually hides a memoised brute-force prime test)

import math

def is_prime(n):
    if n == 2 or n == 3:
        return True
    if n == 1 or n % 2 == 0:
        return False
    else:
        return all(n % i for i in xrange(3, int(1 + math.sqrt(n)), 2))


class Primes(object):

    def __init__(self):
        self.memo = {}

    def __contains__(self, n):
        if n not in self.memo:
            self.memo[n] = is_prime(n)
        return self.memo[n]

That seems to be working so far:

>>> primes = Primes()
>>> 7 in primes
True
>>> 104729 in primes
True
>>> 100 in primes
False
>>> 100 not in primes
True

But it's not playing nicely with argparse:

>>> import argparse as ap
>>> parser = ap.ArgumentParser()
>>> parser.add_argument('prime', type=int, choices=primes, metavar='p')
_StoreAction(option_strings=[], dest='prime', nargs=None, const=None, default=None, type=<type 'int'>, choices=<__main__.Primes object at 0x7f4e21783f10>, help=None, metavar='p')
>>> parser.parse_args(['7'])
Namespace(prime=7)
>>> parser.parse_args(['11'])
Namespace(prime=11)
>>> parser.parse_args(['12'])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python2.7/argparse.py", line 1688, in parse_args
    args, argv = self.parse_known_args(args, namespace)
  File "/usr/lib/python2.7/argparse.py", line 1720, in parse_known_args
    namespace, args = self._parse_known_args(args, namespace)
  File "/usr/lib/python2.7/argparse.py", line 1929, in _parse_known_args
    stop_index = consume_positionals(start_index)
  File "/usr/lib/python2.7/argparse.py", line 1885, in consume_positionals
    take_action(action, args)
  File "/usr/lib/python2.7/argparse.py", line 1778, in take_action
    argument_values = self._get_values(action, argument_strings)
  File "/usr/lib/python2.7/argparse.py", line 2219, in _get_values
    self._check_value(action, value)
  File "/usr/lib/python2.7/argparse.py", line 2267, in _check_value
    tup = value, ', '.join(map(repr, action.choices))
TypeError: argument 2 to map() must support iteration

The docs just say that

Any object that supports the in operator can be passed as the choices value, so dict objects, set objects, custom containers, etc. are all supported.

Obviously I don't want to iterate the infinite "set" of primes. So why the heck is argparse trying to map my primes? Doesn't it just need in and not in?

回答1:

It appears to be a documentation bug. The library as written requires the choices argument to be not just a container but also iterable, It tries to list the available options, which isn't going to work for your case. You could try to hack it by giving it a fake __iter__ that just returns some informational string.

You might also want to submit this as a bug in the Python bug tracker, since the behavior really is contradicted by the docs.

回答2:

The source where you're getting the exception is pretty clear, check it out:

if action.choices is not None and value not in action.choices:
    tup = value, ', '.join(map(repr, action.choices))
    msg = _('invalid choice: %r (choose from %s)') % tup
    raise ArgumentError(action, msg)

The check itself is fine. It borks on trying to print out a helpful error message, where it's trying to give you all the possible choices. I suppose if you define the iterator to return just something that has as its repr the string primes, you might hack it to do the rightish thing.

回答3:

choices are for arguments that you can enumerate all allowed values (a finite (small) set). The docs should be more clear about it.

primes is an infinite set. You could set type parameter to raise ValueError for non-primes.