argparse: Add support for custom argument types.

This commit adds support for optional custom argument type validation to
argparse.ArgumentParser, allowing for shorter argument validation code
for both simple builtins and complex types.

For example, assuming that a particular command line argument must be an
integer, using "parser.add_argument('-a', type=int)" will make sure that
any value passed to that argument that cannot be converted into an
integer will trigger an argument validation error.

Signed-off-by: Alessandro Gatti <[email protected]>

Author: agatti

python-stdlib/argparse/argparse.py

@@ -11,26 +11,36 @@ class _ArgError(BaseException):
 
 
 class _Arg:
-    def __init__(self, names, dest, action, nargs, const, default, help):
+    def __init__(self, names, dest, action, nargs, const, default, help, type):
         self.names = names
         self.dest = dest
         self.action = action
         self.nargs = nargs
         self.const = const
         self.default = default
         self.help = help
+        self.type = type
+
+    def _pop(self, optname, args):
+        arg = args.pop(0)
+        if self.type:
+            try:
+                return self.type(arg)
+            except Exception as e:
+                raise _ArgError("invalid value for %s: %s (%s)" % (optname, arg, str(e)))
+        return arg
 
     def parse(self, optname, args):
         # parse args for this arg
         if self.action == "store":
             if self.nargs is None:
                 if args:
-                    return args.pop(0)
+                    return self._pop(optname, args)
                 else:
                     raise _ArgError("expecting value for %s" % optname)
             elif self.nargs == "?":
                 if args:
-                    return args.pop(0)
+                    return self._pop(optname, args)
                 else:
                     return self.default
             else:
@@ -52,7 +62,7 @@ def parse(self, optname, args):
                         else:
                             break
                     else:
-                        ret.append(args.pop(0))
+                        ret.append(self._pop(optname, args))
                         n -= 1
                 if n > 0:
                     raise _ArgError("expecting value for %s" % optname)
@@ -103,6 +113,9 @@ def add_argument(self, *args, **kwargs):
                 dest = args[0]
             if not args:
                 args = [dest]
+        arg_type = kwargs.get("type", None)
+        if arg_type is not None and not callable(arg_type):
+            raise ValueError("type is not callable")
         list.append(
             _Arg(
                 args,
@@ -112,6 +125,7 @@ def add_argument(self, *args, **kwargs):
                 const,
                 default,
                 kwargs.get("help", ""),
+                arg_type,
             )
         )
 

python-stdlib/argparse/manifest.py

@@ -1,4 +1,4 @@
-metadata(version="0.4.0")
+metadata(version="0.4.1")
 
 # Originally written by Damien George.
 

python-stdlib/argparse/test_argparse.py

@@ -66,3 +66,37 @@
 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"]
+
+
+class CustomArgType:
+    def __init__(self, add):
+        self.add = add
+
+    def __call__(self, value):
+        return int(value) + self.add
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-a", type=int)
+parser.add_argument("-b", type=str)
+parser.add_argument("-c", type=CustomArgType(1))
+try:
+    parser.add_argument("-d", type=())
+    assert False
+except ValueError as e:
+    assert str(e) == "type is not callable"
+parser.add_argument("-d", type=int, nargs="+")
+parser.add_argument("-e", type=CustomArgType(1), nargs="+")
+parser.add_argument("-f", type=CustomArgType(1), nargs="?")
+args = parser.parse_args(["-a", "123"])
+assert args.a == 123
+args = parser.parse_args(["-b", "string"])
+assert args.b == "string"
+args = parser.parse_args(["-c", "123"])
+assert args.c == 124
+args = parser.parse_args(["-d", "123", "124", "125"])
+assert args.d == [123, 124, 125]
+args = parser.parse_args(["-e", "123", "124", "125"])
+assert args.e == [124, 125, 126]
+args = parser.parse_args(["-f", "123"])
+assert args.f == 124