argparse: Implement parse_known_args

stinos · pfalcon · commit 08b522abac4a · 2018-02-02T01:21:32.000+02:00
This is convenient when components need only to parse a subset of an
application's arguments, and can be implemented with minor changes to
_parse_args: basically just add unknown arguments to a list instead of
raising an exception.
diff --git a/argparse/argparse.py b/argparse/argparse.py
@@ -144,25 +144,37 @@ def render_arg(arg):
             print("  %-16s%s" % (', '.join(opt.names) + render_arg(opt), opt.help))
 
     def parse_args(self, args=None):
+        return self._parse_args_impl(args, False)
+
+    def parse_known_args(self, args=None):
+        return self._parse_args_impl(args, True)
+
+    def _parse_args_impl(self, args, return_unknown):
         if args is None:
             args = sys.argv[1:]
         else:
             args = args[:]
         try:
-            return self._parse_args(args)
+            return self._parse_args(args, return_unknown)
         except _ArgError as e:
             self.usage(False)
             print("error:", e)
             sys.exit(2)
 
-    def _parse_args(self, args):
+    def _parse_args(self, args, return_unknown):
         # add optional args with defaults
         arg_dest = []
         arg_vals = []
         for opt in self.opt:
             arg_dest.append(opt.dest)
             arg_vals.append(opt.default)
 
+        # deal with unknown arguments, if needed
+        unknown = []
+        def consume_unknown():
+            while args and not args[0].startswith("-"):
+                unknown.append(args.pop(0))
+
         # parse all args
         parsed_pos = False
         while args or not parsed_pos:
@@ -179,15 +191,26 @@ def _parse_args(self, args):
                         found = True
                         break
                 if not found:
-                    raise _ArgError("unknown option %s" % a)
+                    if return_unknown:
+                        unknown.append(a)
+                        consume_unknown()
+                    else:
+                        raise _ArgError("unknown option %s" % a)
             else:
                 # positional arg
                 if parsed_pos:
-                    raise _ArgError("extra args: %s" % " ".join(args))
+                    if return_unknown:
+                        unknown = unknown + args
+                        break
+                    else:
+                        raise _ArgError("extra args: %s" % " ".join(args))
                 for pos in self.pos:
                     arg_dest.append(pos.dest)
                     arg_vals.append(pos.parse(pos.names[0], args))
                 parsed_pos = True
+                if return_unknown:
+                    consume_unknown()
 
         # build and return named tuple with arg values
-        return namedtuple("args", arg_dest)(*arg_vals)
+        values = namedtuple("args", arg_dest)(*arg_vals)
+        return (values, unknown) if return_unknown else values
diff --git a/argparse/test_argparse.py b/argparse/test_argparse.py
@@ -44,3 +44,25 @@
 assert args.files1 == ["a", "b"] and args.files2 == []
 args = parser.parse_args(["a", "b", "c"])
 assert args.files1 == ["a", "b"] and args.files2 == ["c"]
+
+parser = argparse.ArgumentParser()
+parser.add_argument("a", nargs=2)
+parser.add_argument("-b")
+args, rest = parser.parse_known_args(["a", "b", "-b", "2"])
+assert args.a == ["a", "b"] and args.b == "2"
+assert rest == []
+args, rest = parser.parse_known_args(["-b", "2", "a", "b", "c"])
+assert args.a == ["a", "b"] and args.b == "2"
+assert rest == ["c"]
+args, rest = parser.parse_known_args(["a", "b", "-b", "2", "c"])
+assert args.a == ["a", "b"] and args.b == "2"
+assert rest == ["c"]
+args, rest = parser.parse_known_args(["-b", "2", "a", "b", "-", "c"])
+assert args.a == ["a", "b"] and args.b == "2"
+assert rest == ["-", "c"]
+args, rest = parser.parse_known_args(["a", "b", "-b", "2", "-", "x", "y"])
+assert args.a == ["a", "b"] and args.b == "2"
+assert rest == ["-", "x", "y"]
+args, rest = parser.parse_known_args(["a", "b", "c", "-b", "2", "--x", "5", "1"])
+assert args.a == ["a", "b"] and args.b == "2"
+assert rest == ["c", "--x", "5", "1"]
