Rewrite argument parser

tkf · tkf · commit 82b0b5eec746 · 2018-10-23T21:08:45.000-07:00
diff --git a/julia/pseudo_python_cli.py b/julia/pseudo_python_cli.py
@@ -7,19 +7,43 @@
 
 from __future__ import print_function, absolute_import
 
-import argparse
+from collections import namedtuple
 import code
+import copy
 import runpy
 import sys
 import traceback
 
+try:
+    from types import SimpleNamespace
+except ImportError:
+    from argparse import Namespace as SimpleNamespace
+
+
+ARGUMENT_HELP = """
+positional arguments:
+  script         path to file (default: None)
+  args           arguments passed to program in sys.argv[1:]
+
+optional arguments:
+  -h, --help     show this help message and exit
+  -i             inspect interactively after running script.
+  --version, -V  Print the Python version number and exit.
+                 -VV is not supported.
+  -c COMMAND     Execute the Python code in COMMAND.
+  -m MODULE      Search sys.path for the named MODULE and execute its contents
+                 as the __main__ module.
+"""
+
 
 def python(module, command, script, args, interactive):
     if command:
         sys.argv[0] = "-c"
-    elif script:
-        sys.argv[0] = script
+
+    assert sys.argv
     sys.argv[1:] = args
+    if script:
+        sys.argv[0] = script
 
     banner = ""
     try:
@@ -50,93 +74,221 @@ def python(module, command, script, args, interactive):
     if interactive:
         code.interact(banner=banner, local=scope)
 
+ArgDest = namedtuple("ArgDest", "dest names default")
+Optional = namedtuple("Optional", "name is_long argdest nargs action terminal")
+Result = namedtuple("Result", "option values")
+
+
+class PyArgumentParser(object):
+
+    """
+    `ArgumentParser`-like parser with "terminal option" support.
+
+    Major differences:
+
+    * Formatted help has to be provided to `description`.
+    * Many options for `.add_argument` are not supported.
+    * Especially, there is no positional argument support: all positional
+      arguments go into `ns.args`.
+    * `.add_argument` can take boolean option `terminal` (default: `False`)
+      to stop parsing after consuming the given option.
+    """
+
+    def __init__(self, prog=None, usage="%(prog)s [options] [args]",
+                 description=""):
+        self.prog = sys.argv[0] if prog is None else prog
+        self.usage = usage
+        self.description = description
+
+        self._dests = ["args"]
+        self._argdests = [ArgDest("args", (), [])]
+        self._options = []
+
+        self.add_argument("--help", "-h", "-?", action="store_true")
+
+    def format_usage(self):
+        return "usage: " + self.usage % {"prog": self.prog}
+
+    # Once we drop Python 2, we can do:
+    """
+    def add_argument(self, name, *alt, dest=None, nargs=None, action=None,
+                     default=None, terminal=False):
+    """
+
+    def add_argument(self, name, *alt, **kwargs):
+        return self._add_argument_impl(name, alt, **kwargs)
+
+    def _add_argument_impl(self, name, alt, dest=None, nargs=None, action=None,
+                           default=None, terminal=False):
+        if dest is None:
+            if name.startswith("--"):
+                dest = name[2:]
+            elif not name.startswith("-"):
+                dest = name
+            else:
+                raise ValueError(name)
+
+        if not name.startswith("-"):
+            raise NotImplementedError(
+                "Positional arguments are not supported."
+                " All positional arguments will be stored in `ns.args`.")
+        if terminal and action is not None:
+            raise NotImplementedError("Terminal option has to have argument.")
+
+        if nargs is not None and action is not None:
+            raise TypeError("`nargs` and `action` are mutually exclusive")
+        if action == "store_true":
+            nargs = 0
+        if nargs is None:
+            nargs = 1
+        assert isinstance(nargs, int)
+        assert action in (None, "store_true")
 
-class CustomFormatter(argparse.RawDescriptionHelpFormatter,
-                      argparse.ArgumentDefaultsHelpFormatter):
-    pass
+        assert dest not in self._dests
+        self._dests.append(dest)
 
+        argdest = ArgDest(
+            dest=dest,
+            names=(name,) + alt,
+            default=default,
+        )
+        self._argdests.append(argdest)
 
-def make_parser(description=__doc__):
-    parser = argparse.ArgumentParser(
+        for arg in (name,) + alt:
+            self._options.append(Optional(
+                name=arg,
+                is_long=arg.startswith("--"),
+                argdest=argdest,
+                nargs=nargs,
+                action=action,
+                terminal=terminal,
+            ))
+
+    def parse_args(self, args):
+        ns = SimpleNamespace(**{
+            argdest.dest: copy.copy(argdest.default)
+            for argdest in self._argdests
+        })
+        args_iter = iter(args)
+        self._parse_until_terminal(ns, args_iter)
+        ns.args.extend(args_iter)
+
+        if ns.help:
+            self.print_help()
+            self.exit()
+        del ns.help
+
+        return ns
+
+    def _parse_until_terminal(self, ns, args_iter):
+        seen = set()
+        for a in args_iter:
+
+            results = self._find_matches(a)
+            if not results:
+                ns.args.append(a)
+                break
+
+            for i, res in enumerate(results):
+                dest = res.option.argdest.dest
+                if dest in seen:
+                    self._usage_and_error(
+                        "{} provided more than twice"
+                        .format(" ".join(res.option.argdest.names)))
+                seen.add(dest)
+
+                while len(res.values) < res.option.nargs:
+                    try:
+                        res.values.append(next(args_iter))
+                    except StopIteration:
+                        self.error(self.format_usage())
+
+                if res.option.action == "store_true":
+                    setattr(ns, dest, True)
+                else:
+                    value = res.values
+                    if res.option.nargs == 1:
+                        value, = value
+                    setattr(ns, dest, value)
+
+                if res.option.terminal:
+                    assert i == len(results) - 1
+                    return
+
+    def _find_matches(self, arg):
+        for opt in self._options:
+            if arg == opt.name:
+                return [Result(opt, [])]
+            elif arg.startswith(opt.name):
+                # i.e., len(arg) > len(opt.name):
+                if opt.is_long and arg[len(opt.name)] == "=":
+                    return [Result(opt, [arg[len(opt.name) + 1:]])]
+                elif not opt.is_long:
+                    if opt.nargs > 0:
+                        return [Result(opt, [arg[len(opt.name):]])]
+                    else:
+                        results = [Result(opt, [])]
+                        rest = "-" + arg[len(opt.name):]
+                        results.extend(self._find_matches(rest))
+                        return results
+                        # arg="-ih" -> rest="-h"
+        return []
+
+    def print_usage(self, file=None):
+        print(self.format_usage(), file=file or sys.stdout)
+
+    def print_help(self):
+        self.print_usage()
+        print()
+        print(self.description)
+
+    def exit(self, status=0):
+        sys.exit(status)
+
+    def _usage_and_error(self, message):
+        self.print_usage(sys.stderr)
+        print(file=sys.stderr)
+        self.error(message)
+
+    def error(self, message):
+        print(message, file=sys.stderr)
+        self.exit(2)
+
+
+def make_parser(description=__doc__ + ARGUMENT_HELP):
+    parser = PyArgumentParser(
         prog=None if sys.argv[0] else "python",
         usage="%(prog)s [option] ... [-c cmd | -m mod | script | -] [args]",
-        formatter_class=CustomFormatter,
         description=description)
 
-    parser.add_argument(
-        "-i", dest="interactive", action="store_true",
-        help="""
-        inspect interactively after running script.
-        """)
-    parser.add_argument(
-        "--version", "-V", action="version",
-        version="Python {0}.{1}.{2}".format(*sys.version_info),
-        help="""
-        print the Python version number and exit.
-        -VV is not supported.
-        """)
-
-    group = parser.add_mutually_exclusive_group()
-    group.add_argument(
-        "-c", dest="command",
-        help="""
-        Execute the Python code in COMMAND.
-        """)
-    group.add_argument(
-        "-m", dest="module",
-        help="""
-        Search sys.path for the named MODULE and execute its contents
-        as the __main__ module.
-        """)
-
-    parser.add_argument(
-        "script", nargs="?",
-        help="path to file")
-    parser.add_argument(
-        "args", nargs=argparse.REMAINDER,
-        help="arguments passed to program in sys.argv[1:]")
+    parser.add_argument("-i", dest="interactive", action="store_true")
+    parser.add_argument("--version", "-V", action="store_true")
+    parser.add_argument("-c", dest="command", terminal=True)
+    parser.add_argument("-m", dest="module", terminal=True)
 
     return parser
 
 
 def parse_args_with(parser, args):
-    ns = parser.parse_args(list(preprocess_args(args)))
-    if (ns.command or ns.module) and ns.script:
-        ns.args = [ns.script] + ns.args
-        ns.script = None
-    return ns
+    ns = parser.parse_args(args)
 
+    if ns.command and ns.module:
+        parser.error("-c and -m are mutually exclusive")
+    if ns.version:
+        print("Python {0}.{1}.{2}".format(*sys.version_info))
+        parser.exit()
+    del ns.version
 
-def parse_args(args):
-    return parse_args_with(make_parser(), args)
+    ns.script = None
+    if (not (ns.command or ns.module)) and ns.args:
+        ns.script = ns.args[0]
+        ns.args = ns.args[1:]
 
+    return ns
 
-def preprocess_args(args):
-    """
-    Insert "--" after "[-c cmd | -m mod | script | -]"
-
-    This is required for the following to work:
 
-    >>> ns = parse_args(["-mjson.tool", "-h"])
-    >>> ns.args
-    ['-h']
-    """
-    it = iter(args)
-    for a in it:
-        yield a
-
-        if a in ("-m", "-c"):
-            try:
-                yield next(it)
-            except StopIteration:
-                return
-            yield "--"
-        elif a == "-":
-            yield "--"
-        elif a.startswith("-"):
-            if a[1] in ("m", "c"):
-                yield "--"
-            # otherwise, it's some
+def parse_args(args):
+    return parse_args_with(make_parser(), args)
 
 
 def main(args=None):
diff --git a/julia/python_jl.py b/julia/python_jl.py
@@ -28,7 +28,11 @@
 import os
 import sys
 
-from .pseudo_python_cli import make_parser, parse_args_with
+from .pseudo_python_cli import make_parser, parse_args_with, ARGUMENT_HELP
+
+PYJL_ARGUMENT_HELP = ARGUMENT_HELP + """
+  --julia JULIA  Julia interpreter to be used. (default: julia)
+"""
 
 script_jl = """
 import PyCall
@@ -99,12 +103,8 @@ def parse_pyjl_args(args):
     # parse error right now without initiating Julia interpreter and
     # importing PyCall.jl etc. to get an extra speedup for the
     # abnormal case (including -h/--help and -V/--version).
-    parser = make_parser(description=__doc__)
-    parser.add_argument(
-        "--julia", default="julia",
-        help="""
-        Julia interpreter to be used.
-        """)
+    parser = make_parser(description=__doc__ + PYJL_ARGUMENT_HELP)
+    parser.add_argument("--julia", default="julia")
 
     ns = parse_args_with(parser, args)
     unused_args = list(remove_julia_options(args))
diff --git a/test/test_pseudo_python_cli.py b/test/test_pseudo_python_cli.py
