Bug 1163112 - [mach core] Consolidate functionality between Main._run and Registrar.dispatch, r=gps

--HG--
extra : rebase_source : 37fb8698c388b63b212c4b02b387dd8814eeecde

python/mach/mach/main.py

@@ -16,7 +16,6 @@ import os
 import sys
 import traceback
 import uuid
-import sys
 
 from .base import (
     CommandContext,
@@ -92,13 +91,6 @@ It looks like you passed an unrecognized argument into mach.
 The %s command does not accept the arguments: %s
 '''.lstrip()
 
-INVALID_COMMAND_CONTEXT = r'''
-It looks like you tried to run a mach command from an invalid context. The %s
-command failed to meet the following conditions: %s
-
-Run |mach help| to show a list of all commands available to the current context.
-'''.lstrip()
-
 INVALID_ENTRY_POINT = r'''
 Entry points should return a list of command providers or directories
 containing command providers. The following entry point is invalid:
@@ -158,7 +150,7 @@ class ContextWrapper(object):
         except AttributeError as e:
             try:
                 ret = object.__getattribute__(self, '_handler')(self, key)
-            except AttributeError, TypeError:
+            except (AttributeError, TypeError):
                 # TypeError is in case the handler comes from old code not
                 # taking a key argument.
                 raise e
@@ -426,45 +418,17 @@ To see more help for a specific command, run:
             raise MachError('ArgumentParser result missing mach handler info.')
 
         handler = getattr(args, 'mach_handler')
-        cls = handler.cls
-
-        if handler.pass_context:
-            instance = cls(context)
-        else:
-            instance = cls()
-
-        if handler.conditions:
-            fail_conditions = []
-            for c in handler.conditions:
-                if not c(instance):
-                    fail_conditions.append(c)
-
-            if fail_conditions:
-                print(self._condition_failed_message(handler.name, fail_conditions))
-                return 1
-
-        fn = getattr(instance, handler.method)
 
         try:
-            if args.debug_command:
-                import pdb
-                result = pdb.runcall(fn, **vars(args.command_args))
-            else:
-                result = fn(**vars(args.command_args))
-
-            if not result:
-                result = 0
-
-            assert isinstance(result, (int, long))
-
-            return result
+            return Registrar._run_command_handler(handler, context=context,
+                debug_command=args.debug_command, **vars(args.command_args))
         except KeyboardInterrupt as ki:
             raise ki
         except Exception as e:
             exc_type, exc_value, exc_tb = sys.exc_info()
 
-            # The first frame is us and is never used.
-            stack = traceback.extract_tb(exc_tb)[1:]
+            # The first two frames are us and are never used.
+            stack = traceback.extract_tb(exc_tb)[2:]
 
             # If we have nothing on the stack, the exception was raised as part
             # of calling the @Command method itself. This likely means a
@@ -511,16 +475,6 @@ To see more help for a specific command, run:
         self.logger.log(level, format_str,
             extra={'action': action, 'params': params})
 
-    @classmethod
-    def _condition_failed_message(cls, name, conditions):
-        msg = ['\n']
-        for c in conditions:
-            part = ['  %s' % c.__name__]
-            if c.__doc__ is not None:
-                part.append(c.__doc__)
-            msg.append(' - '.join(part))
-        return INVALID_COMMAND_CONTEXT % (name, '\n'.join(msg))
-
     def _print_error_header(self, argv, fh):
         fh.write('Error running mach:\n\n')
         fh.write('    ')

python/mach/mach/registrar.py

@@ -6,6 +6,13 @@ from __future__ import unicode_literals
 
 from .base import MachError
 
+INVALID_COMMAND_CONTEXT = r'''
+It looks like you tried to run a mach command from an invalid context. The %s
+command failed to meet the following conditions: %s
+
+Run |mach help| to show a list of all commands available to the current context.
+'''.lstrip()
+
 
 class MachRegistrar(object):
     """Container for mach command and config providers."""
@@ -38,15 +45,17 @@ class MachRegistrar(object):
         self.categories[name] = (title, description, priority)
         self.commands_by_category[name] = set()
 
-    def dispatch(self, name, context=None, **args):
-        """Dispatch/run a command.
+    @classmethod
+    def _condition_failed_message(cls, name, conditions):
+        msg = ['\n']
+        for c in conditions:
+            part = ['  %s' % c.__name__]
+            if c.__doc__ is not None:
+                part.append(c.__doc__)
+            msg.append(' - '.join(part))
+        return INVALID_COMMAND_CONTEXT % (name, '\n'.join(msg))
 
-        Commands can use this to call other commands.
-        """
-
-        # TODO The logic in this function overlaps with code in
-        # mach.main.Main._run() and should be consolidated.
-        handler = self.command_handlers[name]
+    def _run_command_handler(self, handler, context=None, debug_command=False, **kwargs):
         cls = handler.cls
 
         if handler.pass_context and not context:
@@ -57,9 +66,49 @@ class MachRegistrar(object):
         else:
             instance = cls()
 
+        if handler.conditions:
+            fail_conditions = []
+            for c in handler.conditions:
+                if not c(instance):
+                    fail_conditions.append(c)
+
+            if fail_conditions:
+                print(self._condition_failed_message(handler.name, fail_conditions))
+                return 1
+
         fn = getattr(instance, handler.method)
 
-        return fn(**args) or 0
+        if debug_command:
+            import pdb
+            result = pdb.runcall(fn, **kwargs)
+        else:
+            result = fn(**kwargs)
+
+        result = result or 0
+        assert isinstance(result, (int, long))
+        return result
+
+    def dispatch(self, name, context=None, argv=None, **kwargs):
+        """Dispatch/run a command.
+
+        Commands can use this to call other commands.
+        """
+        # TODO handler.subcommand_handlers are ignored
+        handler = self.command_handlers[name]
+
+        if handler.parser:
+            parser = handler.parser
+
+            # save and restore existing defaults so **kwargs don't persist across
+            # subsequent invocations of Registrar.dispatch()
+            old_defaults = parser._defaults.copy()
+            parser.set_defaults(**kwargs)
+            kwargs, _ = parser.parse_known_args(argv or [])
+            kwargs = vars(kwargs)
+            parser._defaults = old_defaults
+
+        return self._run_command_handler(handler, context=context, **kwargs)
+
 
 
 Registrar = MachRegistrar()

python/mach/mach/test/test_conditions.py

@@ -8,6 +8,7 @@ import os
 
 from mach.base import MachError
 from mach.main import Mach
+from mach.registrar import Registrar
 from mach.test.common import TestBase
 
 from mozunit import main
@@ -48,14 +49,14 @@ class TestConditions(TestBase):
             result, stdout, stderr = self._run_mach([name])
             self.assertEquals(1, result)
 
-            fail_msg = Mach._condition_failed_message(name, fail_conditions)
+            fail_msg = Registrar._condition_failed_message(name, fail_conditions)
             self.assertEquals(fail_msg.rstrip(), stdout.rstrip())
 
         for name in ('cmd_bar_ctx', 'cmd_foobar_ctx'):
             result, stdout, stderr = self._run_mach([name], _populate_context)
             self.assertEquals(1, result)
 
-            fail_msg = Mach._condition_failed_message(name, fail_conditions)
+            fail_msg = Registrar._condition_failed_message(name, fail_conditions)
             self.assertEquals(fail_msg.rstrip(), stdout.rstrip())
 
     def test_invalid_type(self):