Rework to add tests and make optional

Author: savannahostrowski

Lib/argparse.py

@@ -84,7 +84,6 @@
     'ZERO_OR_MORE',
 ]
 
-import difflib as _difflib
 import os as _os
 import re as _re
 import sys as _sys
@@ -1716,6 +1715,8 @@ class ArgumentParser(_AttributeHolder, _ActionsContainer):
         - allow_abbrev -- Allow long options to be abbreviated unambiguously
         - exit_on_error -- Determines whether or not ArgumentParser exits with
             error info when an error occurs
+        - suggest_on_error - Enables argument choice and subparser name
+            suggestions on user typo
     """
 
     def __init__(self,
@@ -1731,7 +1732,8 @@ def __init__(self,
                  conflict_handler='error',
                  add_help=True,
                  allow_abbrev=True,
-                 exit_on_error=True):
+                 exit_on_error=True,
+                 suggest_on_error=False):
 
         superinit = super(ArgumentParser, self).__init__
         superinit(description=description,
@@ -1751,6 +1753,7 @@ def __init__(self,
         self.add_help = add_help
         self.allow_abbrev = allow_abbrev
         self.exit_on_error = exit_on_error
+        self.suggest_on_error = suggest_on_error
 
         add_group = self.add_argument_group
         self._positionals = add_group(_('positional arguments'))
@@ -2559,22 +2562,23 @@ def _check_value(self, action, value):
 
         # converted value must be one of the choices (if specified)
         if action.choices is not None and value not in action.choices:
-            try:
-                closest_choice = _difflib.get_close_matches(value, action.choices, 1)
-            except TypeError:
-                closest_choice = []
-
             args = {
                 'value': value,
                 'choices': ', '.join(map(repr, action.choices)),
             }
-            if closest_choice:
-                closest_choice = closest_choice[0]
-                args['closest'] = closest_choice
-                msg = _('invalid choice: %(value)r, maybe you meant %(closest)r? '
-                        '(choose from %(choices)s)')
-            else:
-                msg = _('invalid choice: %(value)r (choose from %(choices)s)')
+            msg = _('invalid choice: %(value)r (choose from %(choices)s)')
+
+            if self.suggest_on_error:
+                try:
+                    import difflib as _difflib
+                    closest_choice = _difflib.get_close_matches(value, action.choices, 1)
+                    if closest_choice:
+                        closest_choice = closest_choice[0]
+                        args['closest'] = closest_choice
+                        msg = _('invalid choice: %(value)r, maybe you meant %(closest)r? '
+                                '(choose from %(choices)s)')
+                except TypeError:
+                    closest_choice = []
 
             raise ArgumentError(action, msg % args)
 

Lib/test/test_argparse.py

@@ -2172,6 +2172,53 @@ class TestNegativeNumber(ParserTestCase):
         ('--float -.5_000', NS(int=None, float=-0.5)),
     ]
 
+class TestArgumentAndSubparserSuggestions(TestCase):
+    """Test error handling and suggestion when a user makes a typo"""
+
+    def test_wrong_argument_error_with_suggestions(self):
+        parser = ErrorRaisingArgumentParser(suggest_on_error=True)
+        parser.add_argument('foo', choices=['bar', 'baz'])
+        with self.assertRaises(ArgumentParserError) as excinfo:
+            parser.parse_args(('bazz',))
+        self.assertIn(
+            "maybe you meant 'baz'? (choose from 'bar', 'baz')",
+            excinfo.exception.stderr,
+        )
+
+    def test_wrong_argument_error_no_suggestions(self):
+        parser = ErrorRaisingArgumentParser(suggest_on_error=False)
+        parser.add_argument('foo', choices=['bar', 'baz'])
+        with self.assertRaises(ArgumentParserError) as excinfo:
+            parser.parse_args(('bazz',))
+        self.assertIn(
+            "invalid choice: 'bazz' (choose from 'bar', 'baz')",
+            excinfo.exception.stderr,
+        )
+
+    def test_wrong_argument_subparsers_with_suggestions(self):
+        parser = ErrorRaisingArgumentParser(suggest_on_error=True)
+        subparsers = parser.add_subparsers(required=True)
+        subparsers.add_parser('foo')
+        subparsers.add_parser('bar')
+        with self.assertRaises(ArgumentParserError) as excinfo:
+            parser.parse_args(('baz',))
+        self.assertIn(
+            "maybe you meant 'bar'? (choose from 'foo', 'bar')",
+            excinfo.exception.stderr,
+        )
+
+    def test_wrong_argument_subparsers_no_suggestions(self):
+        parser = ErrorRaisingArgumentParser(suggest_on_error=False)
+        subparsers = parser.add_subparsers(required=True)
+        subparsers.add_parser('foo')
+        subparsers.add_parser('bar')
+        with self.assertRaises(ArgumentParserError) as excinfo:
+            parser.parse_args(('baz',))
+        self.assertIn(
+            "invalid choice: 'baz' (choose from 'foo', 'bar')",
+            excinfo.exception.stderr,
+        )
+
 class TestInvalidAction(TestCase):
     """Test invalid user defined Action"""
 
@@ -2390,18 +2437,6 @@ def test_required_subparsers_no_destination_error(self):
             'error: the following arguments are required: {foo,bar}\n$'
         )
 
-    def test_wrong_argument_subparsers_no_destination_error(self):
-        parser = ErrorRaisingArgumentParser()
-        subparsers = parser.add_subparsers(required=True)
-        subparsers.add_parser('foo')
-        subparsers.add_parser('bar')
-        with self.assertRaises(ArgumentParserError) as excinfo:
-            parser.parse_args(('baz',))
-        self.assertIn(
-            "error: argument {foo,bar}: invalid choice: 'baz', maybe you meant 'bar'? (choose from 'foo', 'bar')",
-            excinfo.exception.stderr,
-        )
-
     def test_optional_subparsers(self):
         parser = ErrorRaisingArgumentParser()
         subparsers = parser.add_subparsers(dest='command', required=False)
@@ -2726,7 +2761,7 @@ def test_single_parent_mutex(self):
         parser = ErrorRaisingArgumentParser(parents=[self.ab_mutex_parent])
         self._test_mutex_ab(parser.parse_args)
 
-    def test_single_granparent_mutex(self):
+    def test_single_grandparent_mutex(self):
         parents = [self.ab_mutex_parent]
         parser = ErrorRaisingArgumentParser(add_help=False, parents=parents)
         parser = ErrorRaisingArgumentParser(parents=[parser])