Bug 901972 - Add ability to filter mach commands at runtime, r=gps

python/mach/README.rst

@@ -55,6 +55,51 @@ to the decorators are being used as arguments to
 The Python modules defining mach commands do not need to live inside the
 main mach source tree.
 
+Conditionally Filtering Commands
+--------------------------------
+
+Sometimes it might only make sense to run a command given a certain
+context. For example, running tests only makes sense if the product
+they are testing has been built, and said build is available. To make
+sure a command is only runnable from within a correct context, you can
+define a series of conditions on the *Command* decorator.
+
+A condition is simply a function that takes an instance of the
+*CommandProvider* class as an argument, and returns True or False. If
+any of the conditions defined on a command return False, the command
+will not be runnable. The doc string of a condition function is used in
+error messages, to explain why the command cannot currently be run.
+
+Here is an example:
+
+    from mach.decorators import (
+        CommandProvider,
+        Command,
+    )
+
+    def build_available(cls):
+        """The build needs to be available."""
+        return cls.build_path is not None
+
+    @CommandProvider
+    class MyClass(MachCommandBase):
+        def __init__(self, build_path=None):
+            self.build_path = build_path
+
+        @Command('run_tests', conditions=[build_available])
+        def run_tests(self, force=False):
+            # Do stuff here.
+
+It is important to make sure that any state needed by the condition is
+available to instances of the command provider.
+
+By default all commands without any conditions applied will be runnable,
+but it is possible to change this behaviour by setting *require_conditions*
+to True:
+
+    m = mach.main.Mach()
+    m.require_conditions = True
+
 Minimizing Code in Mach
 -----------------------
 

python/mach/mach/base.py

@@ -77,13 +77,18 @@ class MethodHandler(object):
         # Whether to allow all arguments from the parser.
         'allow_all_arguments',
 
+        # Functions used to 'skip' commands if they don't meet the conditions
+        # in a given context.
+        'conditions',
+
         # Arguments added to this command's parser. This is a 2-tuple of
         # positional and named arguments, respectively.
         'arguments',
     )
 
     def __init__(self, cls, method, name, category=None, description=None,
-        allow_all_arguments=False, arguments=None, pass_context=False):
+        allow_all_arguments=False, conditions=None, arguments=None,
+        pass_context=False):
 
         self.cls = cls
         self.method = method
@@ -91,6 +96,7 @@ class MethodHandler(object):
         self.category = category
         self.description = description
         self.allow_all_arguments = allow_all_arguments
+        self.conditions = conditions or []
         self.arguments = arguments or []
         self.pass_context = pass_context
 

python/mach/mach/decorators.py

@@ -4,6 +4,7 @@
 
 from __future__ import unicode_literals
 
+import collections
 import inspect
 import types
 
@@ -55,17 +56,35 @@ def CommandProvider(cls):
         if not isinstance(value, types.FunctionType):
             continue
 
-        command_name, category, description, allow_all = getattr(value,
-            '_mach_command', (None, None, None, None))
+        command_name, category, description, allow_all, conditions = getattr(
+            value, '_mach_command', (None, None, None, None, None))
 
         if command_name is None:
             continue
 
+        if conditions is None and Registrar.require_conditions:
+            continue
+
+        msg = 'Mach command \'%s\' implemented incorrectly. ' + \
+              'Conditions argument must take a list ' + \
+              'of functions. Found %s instead.'
+
+        conditions = conditions or []
+        if not isinstance(conditions, collections.Iterable):
+            msg = msg % (command_name, type(conditions))
+            raise MachError(msg)
+
+        for c in conditions:
+            if not hasattr(c, '__call__'):
+                msg = msg % (command_name, type(c))
+                raise MachError(msg)
+
         arguments = getattr(value, '_mach_command_args', None)
 
         handler = MethodHandler(cls, attr, command_name, category=category,
             description=description, allow_all_arguments=allow_all,
-            arguments=arguments, pass_context=pass_context)
+            conditions=conditions, arguments=arguments,
+            pass_context=pass_context)
 
         Registrar.register_command_handler(handler)
 
@@ -93,15 +112,16 @@ class Command(object):
             pass
     """
     def __init__(self, name, category=None, description=None,
-        allow_all_args=False):
+        allow_all_args=False, conditions=None):
         self._name = name
         self._category = category
         self._description = description
         self._allow_all_args = allow_all_args
+        self._conditions = conditions
 
     def __call__(self, func):
         func._mach_command = (self._name, self._category, self._description,
-            self._allow_all_args)
+            self._allow_all_args, self._conditions)
 
         return func
 

python/mach/mach/dispatcher.py

@@ -54,7 +54,7 @@ class CommandAction(argparse.Action):
     For more, read the docs in __call__.
     """
     def __init__(self, option_strings, dest, required=True, default=None,
-        registrar=None):
+        registrar=None, context=None):
         # A proper API would have **kwargs here. However, since we are a little
         # hacky, we intentionally omit it as a way of detecting potentially
         # breaking changes with argparse's implementation.
@@ -65,6 +65,7 @@ class CommandAction(argparse.Action):
             help=argparse.SUPPRESS, nargs=argparse.REMAINDER)
 
         self._mach_registrar = registrar
+        self._context = context
 
     def __call__(self, parser, namespace, values, option_string=None):
         """This is called when the ArgumentParser has reached our arguments.
@@ -148,17 +149,33 @@ class CommandAction(argparse.Action):
         cats = [(k, v[2]) for k, v in r.categories.items()]
         sorted_cats = sorted(cats, key=itemgetter(1), reverse=True)
         for category, priority in sorted_cats:
-            if not r.commands_by_category[category]:
-                continue
-
-            title, description, _priority = r.categories[category]
-
-            group = parser.add_argument_group(title, description)
+            group = None
 
             for command in sorted(r.commands_by_category[category]):
                 handler = r.command_handlers[command]
-                description = handler.description
 
+                # Instantiate a handler class to see if it should be filtered
+                # out for the current context or not. Condition functions can be
+                # applied to the command's decorator.
+                if handler.conditions:
+                    if handler.pass_context:
+                        instance = handler.cls(self._context)
+                    else:
+                        instance = handler.cls()
+
+                    is_filtered = False
+                    for c in handler.conditions:
+                        if not c(instance):
+                            is_filtered = True
+                            break
+                    if is_filtered:
+                        continue
+
+                if group is None:
+                    title, description, _priority = r.categories[category]
+                    group = parser.add_argument_group(title, description)
+
+                description = handler.description
                 group.add_argument(command, help=description,
                     action='store_true')
 

python/mach/mach/main.py

@@ -87,6 +87,13 @@ 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()
+
 
 class ArgumentParser(argparse.ArgumentParser):
     """Custom implementation argument parser to make things look pretty."""
@@ -140,6 +147,10 @@ class Mach(object):
             as its single argument right before command dispatch. This allows
             modification of the context instance and thus passing of
             arbitrary data to command handlers.
+
+        require_conditions -- If True, commands that do not have any condition
+            functions applied will be skipped. Defaults to False.
+
     """
 
     USAGE = """%(prog)s [global arguments] command [command arguments]
@@ -209,6 +220,14 @@ To see more help for a specific command, run:
 
         Registrar.register_category(name, title, description, priority)
 
+    @property
+    def require_conditions(self):
+        return Registrar.require_conditions
+
+    @require_conditions.setter
+    def require_conditions(self, value):
+        Registrar.require_conditions = value
+
     def run(self, argv, stdin=None, stdout=None, stderr=None):
         """Runs mach with arguments provided from the command line.
 
@@ -270,7 +289,14 @@ To see more help for a specific command, run:
             sys.stderr = orig_stderr
 
     def _run(self, argv):
-        parser = self.get_argument_parser()
+        context = CommandContext(topdir=self.cwd, cwd=self.cwd,
+            settings=self.settings, log_manager=self.log_manager,
+            commands=Registrar)
+
+        if self.populate_context_handler:
+            self.populate_context_handler(context)
+
+        parser = self.get_argument_parser(context)
 
         if not len(argv):
             # We don't register the usage until here because if it is globally
@@ -314,13 +340,6 @@ To see more help for a specific command, run:
         if not hasattr(args, 'mach_handler'):
             raise MachError('ArgumentParser result missing mach handler info.')
 
-        context = CommandContext(topdir=self.cwd, cwd=self.cwd,
-            settings=self.settings, log_manager=self.log_manager,
-            commands=Registrar)
-
-        if self.populate_context_handler:
-            self.populate_context_handler(context)
-
         handler = getattr(args, 'mach_handler')
         cls = handler.cls
 
@@ -329,6 +348,16 @@ To see more help for a specific command, run:
         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:
@@ -393,6 +422,16 @@ 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('    ')
@@ -448,7 +487,7 @@ To see more help for a specific command, run:
 
         return os.path.exists(p)
 
-    def get_argument_parser(self):
+    def get_argument_parser(self, context):
         """Returns an argument parser for the command-line interface."""
 
         parser = ArgumentParser(add_help=False,
@@ -476,7 +515,7 @@ To see more help for a specific command, run:
         # We need to be last because CommandAction swallows all remaining
         # arguments and argparse parses arguments in the order they were added.
         parser.add_argument('command', action=CommandAction,
-            registrar=Registrar)
+            registrar=Registrar, context=context)
 
         return parser
 

python/mach/mach/registrar.py

@@ -15,6 +15,7 @@ class MachRegistrar(object):
         self.commands_by_category = {}
         self.settings_providers = set()
         self.categories = {}
+        self.require_conditions = False
 
     def register_command_handler(self, handler):
         name = handler.name

python/mach/mach/test/common.py

@@ -4,26 +4,33 @@
 
 from __future__ import unicode_literals
 
-import time
+from StringIO import StringIO
+import os
+import unittest
 
-from mach.base import (
-    CommandArgument,
-    CommandProvider,
-    Command,
-)
+from mach.main import Mach
+from mach.base import CommandContext
 
-import mach.test.common2 as common2
+here = os.path.abspath(os.path.dirname(__file__))
 
+class TestBase(unittest.TestCase):
+    provider_dir = os.path.join(here, 'providers')
 
-@CommandProvider
-class TestCommandProvider(object):
-    @Command('throw')
-    @CommandArgument('--message', '-m', default='General Error')
-    def throw(self, message):
-        raise Exception(message)
+    def _run_mach(self, args, provider_file, context_handler=None):
+        m = Mach(os.getcwd())
+        m.define_category('testing', 'Mach unittest', 'Testing for mach core', 10)
+        m.populate_context_handler = context_handler
 
-    @Command('throw_deep')
-    @CommandArgument('--message', '-m', default='General Error')
-    def throw_deep(self, message):
-        common2.throw_deep(message)
+        m.load_commands_from_file(os.path.join(self.provider_dir, provider_file))
 
+        stdout = StringIO()
+        stderr = StringIO()
+        stdout.encoding = 'UTF-8'
+        stderr.encoding = 'UTF-8'
+
+        try:
+            result = m.run(args, stdout=stdout, stderr=stderr)
+        except SystemExit:
+            result = None
+
+        return (result, stdout.getvalue(), stderr.getvalue())

python/mach/mach/test/providers/conditions.py

@@ -0,0 +1,53 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from __future__ import unicode_literals
+
+from mach.decorators import (
+    CommandProvider,
+    Command,
+)
+
+def is_foo(cls):
+    """Foo must be true"""
+    return cls.foo
+
+def is_bar(cls):
+    """Bar must be true"""
+    return cls.bar
+
+@CommandProvider
+class ConditionsProvider(object):
+    foo = True
+    bar = False
+
+    @Command('cmd_foo', category='testing', conditions=[is_foo])
+    def run_foo(self):
+        pass
+
+    @Command('cmd_bar', category='testing', conditions=[is_bar])
+    def run_bar(self):
+        pass
+
+    @Command('cmd_foobar', category='testing', conditions=[is_foo, is_bar])
+    def run_foobar(self):
+        pass
+
+@CommandProvider
+class ConditionsContextProvider(object):
+    def __init__(self, context):
+        self.foo = context.foo
+        self.bar = context.bar
+
+    @Command('cmd_foo_ctx', category='testing', conditions=[is_foo])
+    def run_foo(self):
+        pass
+
+    @Command('cmd_bar_ctx', category='testing', conditions=[is_bar])
+    def run_bar(self):
+        pass
+
+    @Command('cmd_foobar_ctx', category='testing', conditions=[is_foo, is_bar])
+    def run_foobar(self):
+        pass

python/mach/mach/test/providers/conditions_invalid.py

@@ -0,0 +1,16 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from __future__ import unicode_literals
+
+from mach.decorators import (
+    CommandProvider,
+    Command,
+)
+
+@CommandProvider
+class ConditionsProvider(object):
+    @Command('cmd_foo', category='testing', conditions=["invalid"])
+    def run_foo(self):
+        pass

python/mach/mach/test/providers/throw.py

@@ -0,0 +1,29 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from __future__ import unicode_literals
+
+import time
+
+from mach.decorators import (
+    CommandArgument,
+    CommandProvider,
+    Command,
+)
+
+from mach.test.providers import throw2
+
+
+@CommandProvider
+class TestCommandProvider(object):
+    @Command('throw')
+    @CommandArgument('--message', '-m', default='General Error')
+    def throw(self, message):
+        raise Exception(message)
+
+    @Command('throw_deep')
+    @CommandArgument('--message', '-m', default='General Error')
+    def throw_deep(self, message):
+        throw2.throw_deep(message)
+

python/mach/mach/test/test_conditions.py

@@ -0,0 +1,73 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from __future__ import unicode_literals
+
+import os
+import unittest
+
+from StringIO import StringIO
+
+from mach.base import MachError
+from mach.main import Mach
+from mach.test.common import TestBase
+
+def _populate_context(context):
+    context.foo = True
+    context.bar = False
+
+class TestConditions(TestBase):
+    """Tests for conditionally filtering commands."""
+
+    def _run_mach(self, args, context_handler=None):
+        return TestBase._run_mach(self, args, 'conditions.py',
+                                  context_handler=context_handler)
+
+
+    def test_conditions_pass(self):
+        """Test that a command which passes its conditions is runnable."""
+
+        self.assertEquals((0, '', ''), self._run_mach(['cmd_foo']))
+        self.assertEquals((0, '', ''), self._run_mach(['cmd_foo_ctx'], _populate_context))
+
+    def test_invalid_context_message(self):
+        """Test that commands which do not pass all their conditions
+        print the proper failure message."""
+
+        def is_bar():
+            """Bar must be true"""
+        fail_conditions = [is_bar]
+
+        for name in ('cmd_bar', 'cmd_foobar'):
+            result, stdout, stderr = self._run_mach([name])
+            self.assertEquals(1, result)
+
+            fail_msg = Mach._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)
+            self.assertEquals(fail_msg.rstrip(), stdout.rstrip())
+
+    def test_invalid_type(self):
+        """Test that a condition which is not callable raises an exception."""
+
+        m = Mach(os.getcwd())
+        m.define_category('testing', 'Mach unittest', 'Testing for mach core', 10)
+        self.assertRaises(MachError, m.load_commands_from_file,
+                os.path.join(self.provider_dir, 'conditions_invalid.py'))
+
+    def test_help_message(self):
+        """Test that commands that are not runnable do not show up in help."""
+
+        result, stdout, stderr = self._run_mach(['help'], _populate_context)
+        self.assertIn('cmd_foo', stdout)
+        self.assertNotIn('cmd_bar', stdout)
+        self.assertNotIn('cmd_foobar', stdout)
+        self.assertIn('cmd_foo_ctx', stdout)
+        self.assertNotIn('cmd_bar_ctx', stdout)
+        self.assertNotIn('cmd_foobar_ctx', stdout)

python/mach/mach/test/test_error_output.py

@@ -11,26 +11,17 @@ import unittest
 
 from StringIO import StringIO
 
-import mach.main
+from mach.main import (
+    COMMAND_ERROR,
+    MODULE_ERROR
+)
+from mach.test.common import TestBase
 
 
-class TestErrorOutput(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls):
-        common_path = os.path.join(os.path.dirname(__file__), 'common.py')
-        imp.load_source('mach.commands.error_output_test', common_path)
+class TestErrorOutput(TestBase):
 
     def _run_mach(self, args):
-        m = mach.main.Mach(os.getcwd())
-
-        stdout = StringIO()
-        stderr = StringIO()
-        stdout.encoding = 'UTF-8'
-        stderr.encoding = 'UTF-8'
-
-        result = m.run(args, stdout=stdout, stderr=stderr)
-
-        return (result, stdout.getvalue(), stderr.getvalue())
+        return TestBase._run_mach(self, args, 'throw.py')
 
     def test_command_error(self):
         result, stdout, stderr = self._run_mach(['throw', '--message',