gh-132558: allow choices to be specified as strings in presence of a type argument

Lib/argparse.py

@@ -1777,6 +1777,8 @@ class ArgumentParser(_AttributeHolder, _ActionsContainer):
             error info when an error occurs
         - suggest_on_error - Enables suggestions for mistyped argument choices
             and subparser names. (default: ``False``)
+        - convert_choices - Runs the ``choices`` through the ``type`` function
+            during checking. (default: ``False``)
     """
 
     def __init__(self,
@@ -1793,7 +1795,8 @@ def __init__(self,
                  add_help=True,
                  allow_abbrev=True,
                  exit_on_error=True,
-                 suggest_on_error=False):
+                 suggest_on_error=False,
+                 convert_choices=False):
 
         superinit = super(ArgumentParser, self).__init__
         superinit(description=description,
@@ -1810,6 +1813,7 @@ def __init__(self,
         self.allow_abbrev = allow_abbrev
         self.exit_on_error = exit_on_error
         self.suggest_on_error = suggest_on_error
+        self.convert_choices = convert_choices
 
         add_group = self.add_argument_group
         self._positionals = add_group(_('positional arguments'))
@@ -2524,7 +2528,7 @@ def _get_values(self, action, arg_strings):
         elif len(arg_strings) == 1 and action.nargs in [None, OPTIONAL]:
             arg_string, = arg_strings
             value = self._get_value(action, arg_string)
-            self._check_value(action, value)
+            self._check_value(action, value, arg_string)
 
         # REMAINDER arguments convert all values, checking none
         elif action.nargs == REMAINDER:
@@ -2533,7 +2537,7 @@ def _get_values(self, action, arg_strings):
         # PARSER arguments convert all values, but check only the first
         elif action.nargs == PARSER:
             value = [self._get_value(action, v) for v in arg_strings]
-            self._check_value(action, value[0])
+            self._check_value(action, value[0], arg_strings[0])
 
         # SUPPRESS argument does not put anything in the namespace
         elif action.nargs == SUPPRESS:
@@ -2542,8 +2546,8 @@ def _get_values(self, action, arg_strings):
         # all other types of nargs produce a list
         else:
             value = [self._get_value(action, v) for v in arg_strings]
-            for v in value:
-                self._check_value(action, v)
+            for v, s in zip(value, arg_strings):
+                self._check_value(action, v, s)
 
         # return the converted value
         return value
@@ -2572,24 +2576,38 @@ def _get_value(self, action, arg_string):
         # return the converted value
         return result
 
-    def _check_value(self, action, value):
+    def _check_value(self, action, value, arg_string=None):
         # converted value must be one of the choices (if specified)
         choices = action.choices
         if choices is None:
             return
 
+        if arg_string is None:
+            arg_string = value
+
         if isinstance(choices, str):
             choices = iter(choices)
 
-        if value not in choices:
-            args = {'value': str(value),
+        typed_choices = []
+        if (self.convert_choices and
+            action.type and
+            all(isinstance(choice, str) for choice in choices)
+        ):
+            try:
+                typed_choices = [action.type(v) for v in choices]
+            except Exception:
+                # We use a blanket catch here, because type is user provided.
+                pass
+
+        if value not in choices and value not in typed_choices:
+            args = {'value': arg_string,
                     'choices': ', '.join(map(str, action.choices))}
             msg = _('invalid choice: %(value)r (choose from %(choices)s)')
 
-            if self.suggest_on_error and isinstance(value, str):
+            if self.suggest_on_error:
                 if all(isinstance(choice, str) for choice in action.choices):
                     import difflib
-                    suggestions = difflib.get_close_matches(value, action.choices, 1)
+                    suggestions = difflib.get_close_matches(arg_string, action.choices, 1)
                     if suggestions:
                         args['closest'] = suggestions[0]
                         msg = _('invalid choice: %(value)r, maybe you meant %(closest)r? '

Lib/test/test_argparse.py

@@ -1934,6 +1934,46 @@ def setUp(self):
         ('-x - -', NS(x=eq_bstdin, spam=eq_bstdin)),
     ]
 
+class TestChoices(ParserTestCase):
+    """Test the original behavior"""
+    def to_dow(arg):
+        days = ["mo", "tu", "we", "th", "fr", "sa", "su"]
+        if arg in days:
+            return days.index(arg) + 1
+        else:
+            return None
+
+    argument_signatures = [
+        Sig('when',
+            type=to_dow, choices=[1, 2, 3, 4, 5, 6, 7],
+        )
+    ]
+    failures = ['now', '1']
+    successes = [
+        ('mo', NS(when=1)),
+        ('su', NS(when=7)),
+    ]
+
+class TestTypedChoices(TestChoices):
+    """Test a set of string choices that convert to weekdays"""
+
+    parser_signature = Sig(convert_choices=True)
+    argument_signatures = [
+        Sig('when',
+            type=TestChoices.to_dow, choices=["mo", "tu", "we" , "th", "fr", "sa", "su"],
+        )
+    ]
+
+class TestTypedChoicesNoFlag(TestChoices):
+    """Without the feature flag we fail"""
+    argument_signatures = [
+        Sig('when',
+            type=TestChoices.to_dow, choices=["mo", "tu", "we" , "th", "fr", "sa", "su"],
+        )
+    ]
+    failures = ['mo']
+    successes = []
+
 
 class WFile(object):
     seen = set()
@@ -5468,6 +5508,39 @@ def custom_type(string):
     version = ''
 
 
+class TestHelpTypedChoices(HelpTestCase):
+    from datetime import date, timedelta
+    def to_date(arg):
+        if arg == "today":
+            return date.today()
+        elif arg == "tomorrow":
+            return date.today() + timedelta(days=1).date()
+        else:
+            return None
+
+    parser_signature = Sig(prog='PROG', convert_choices=True)
+    argument_signatures = [
+        Sig('when',
+            type=to_date,
+            choices=["today", "tomorrow"]
+        ),
+    ]
+
+    usage = '''\
+usage: PROG [-h] {today,tomorrow}
+        '''
+    help = usage + '''\
+
+positional arguments:
+  {today,tomorrow}
+
+options:
+  -h, --help        show this help message and exit
+        '''
+    version = ''
+
+
+
 class TestHelpUsageLongSubparserCommand(TestCase):
     """Test that subparser commands are formatted correctly in help"""
     maxDiff = None