Improve shell-completion wrappers

mike/arguments.py

@@ -0,0 +1,33 @@
+from argparse import *
+
+_ArgumentParser = ArgumentParser
+_Action = Action
+
+
+# Add some simple wrappers to make it easier to specify shell-completion
+# behaviors.
+
+def _add_complete(argument, complete):
+    if complete is not None:
+        argument.complete = complete
+    return argument
+
+
+class Action(_Action):
+    def __init__(self, *args, complete=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        _add_complete(self, complete)
+
+
+class ArgumentParser(_ArgumentParser):
+    @staticmethod
+    def _wrap_complete(action):
+        def wrapper(*args, complete=None, **kwargs):
+            return _add_complete(action(*args, **kwargs), complete)
+
+        return wrapper
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        for k, v in self._registries['action'].items():
+            self._registries['action'][k] = self._wrap_complete(v)

mike/driver.py

@@ -1,8 +1,7 @@
-import argparse
 import os
 import sys
 
-from . import commands
+from . import arguments, commands
 from . import git_utils
 from . import mkdocs_utils
 from .app_version import version as app_version
@@ -56,23 +55,6 @@
 """
 
 
-class CompletingArgumentParser(argparse.ArgumentParser):
-    @staticmethod
-    def _wrap_complete(action):
-        def wrapper(*args, complete=None, **kwargs):
-            argument = action(*args, **kwargs)
-            if complete is not None:
-                argument.complete = complete
-            return argument
-
-        return wrapper
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        for k, v in self._registries['action'].items():
-            self._registries['action'][k] = self._wrap_complete(v)
-
-
 def add_git_arguments(parser, *, commit=True, prefix=True):
     # Add this whenever we add git arguments since we pull the remote and
     # branch from mkdocs.yml.
@@ -245,7 +227,7 @@ def generate_completion(parser, args):
 
 
 def main():
-    parser = CompletingArgumentParser(prog='mike', description=description)
+    parser = arguments.ArgumentParser(prog='mike', description=description)
     subparsers = parser.add_subparsers(metavar='COMMAND')
     subparsers.required = True
 
@@ -345,7 +327,7 @@ def main():
         'help', help='show this help message and exit', add_help=False
     )
     help_p.set_defaults(func=help)
-    help_p.add_argument('subcommand', metavar='CMD', nargs=argparse.REMAINDER,
+    help_p.add_argument('subcommand', metavar='CMD', nargs=arguments.REMAINDER,
                         help='subcommand to request help for')
 
     completion_p = subparsers.add_parser(

test/unit/test_arguments.py

@@ -0,0 +1,27 @@
+from unittest import TestCase
+
+from mike import arguments
+
+
+class TestParser(TestCase):
+    def test_complete(self):
+        p = arguments.ArgumentParser()
+        arg = p.add_argument('--arg', complete='file')
+        self.assertEqual(arg.complete, 'file')
+
+    def test_complete_group(self):
+        p = arguments.ArgumentParser()
+        g = p.add_argument_group()
+        arg = g.add_argument('--arg', complete='file')
+        self.assertEqual(arg.complete, 'file')
+
+    def test_complete_action(self):
+        class MyAction(arguments.Action):
+            def __call__(self, parser, namespace, values, option_string=None):
+                setattr(namespace, self.dest, values.upper())
+
+        p = arguments.ArgumentParser()
+        arg = p.add_argument('--arg', action=MyAction, complete='file')
+        self.assertEqual(arg.complete, 'file')
+        self.assertEqual(p.parse_args(['--arg=foo']),
+                         arguments.Namespace(arg='FOO'))