1# This Source Code Form is subject to the terms of the Mozilla Public 2# License, v. 2.0. If a copy of the MPL was not distributed with this 3# file, You can obtain one at http://mozilla.org/MPL/2.0/. 4 5from __future__ import absolute_import, print_function, unicode_literals 6 7import time 8 9import six 10 11from .base import MachError 12 13INVALID_COMMAND_CONTEXT = r""" 14It looks like you tried to run a mach command from an invalid context. The %s 15command failed to meet the following conditions: %s 16 17Run |mach help| to show a list of all commands available to the current context. 18""".lstrip() 19 20 21class MachRegistrar(object): 22 """Container for mach command and config providers.""" 23 24 def __init__(self): 25 self.command_handlers = {} 26 self.commands_by_category = {} 27 self.settings_providers = set() 28 self.categories = {} 29 self.require_conditions = False 30 self.command_depth = 0 31 32 def register_command_handler(self, handler): 33 name = handler.name 34 35 if not handler.category: 36 raise MachError( 37 "Cannot register a mach command without a " "category: %s" % name 38 ) 39 40 if handler.category not in self.categories: 41 raise MachError( 42 "Cannot register a command to an undefined " 43 "category: %s -> %s" % (name, handler.category) 44 ) 45 46 self.command_handlers[name] = handler 47 self.commands_by_category[handler.category].add(name) 48 49 def register_settings_provider(self, cls): 50 self.settings_providers.add(cls) 51 52 def register_category(self, name, title, description, priority=50): 53 self.categories[name] = (title, description, priority) 54 self.commands_by_category[name] = set() 55 56 @classmethod 57 def _condition_failed_message(cls, name, conditions): 58 msg = ["\n"] 59 for c in conditions: 60 part = [" %s" % getattr(c, "__name__", c)] 61 if c.__doc__ is not None: 62 part.append(c.__doc__) 63 msg.append(" - ".join(part)) 64 return INVALID_COMMAND_CONTEXT % (name, "\n".join(msg)) 65 66 @classmethod 67 def _instance(_, handler, context, **kwargs): 68 if context is None: 69 raise ValueError("Expected a non-None context.") 70 71 prerun = getattr(context, "pre_dispatch_handler", None) 72 if prerun: 73 prerun(context, handler, args=kwargs) 74 75 context.handler = handler 76 return handler.create_instance(context, handler.virtualenv_name) 77 78 @classmethod 79 def _fail_conditions(_, handler, instance): 80 fail_conditions = [] 81 if handler.conditions: 82 for c in handler.conditions: 83 if not c(instance): 84 fail_conditions.append(c) 85 86 return fail_conditions 87 88 def _run_command_handler(self, handler, context, debug_command=False, **kwargs): 89 instance = MachRegistrar._instance(handler, context, **kwargs) 90 fail_conditions = MachRegistrar._fail_conditions(handler, instance) 91 if fail_conditions: 92 print( 93 MachRegistrar._condition_failed_message(handler.name, fail_conditions) 94 ) 95 return 1 96 97 self.command_depth += 1 98 fn = getattr(instance, handler.method) 99 100 start_time = time.time() 101 102 if debug_command: 103 import pdb 104 105 result = pdb.runcall(fn, instance, **kwargs) 106 else: 107 result = fn(instance, **kwargs) 108 109 end_time = time.time() 110 111 result = result or 0 112 assert isinstance(result, six.integer_types) 113 114 if not debug_command: 115 postrun = getattr(context, "post_dispatch_handler", None) 116 if postrun: 117 postrun( 118 context, 119 handler, 120 instance, 121 not result, 122 start_time, 123 end_time, 124 self.command_depth, 125 args=kwargs, 126 ) 127 self.command_depth -= 1 128 129 return result 130 131 def dispatch(self, name, context, argv=None, subcommand=None, **kwargs): 132 """Dispatch/run a command. 133 134 Commands can use this to call other commands. 135 """ 136 handler = self.command_handlers[name] 137 138 if subcommand: 139 handler = handler.subcommand_handlers[subcommand] 140 141 if handler.parser: 142 parser = handler.parser 143 144 # save and restore existing defaults so **kwargs don't persist across 145 # subsequent invocations of Registrar.dispatch() 146 old_defaults = parser._defaults.copy() 147 parser.set_defaults(**kwargs) 148 kwargs, unknown = parser.parse_known_args(argv or []) 149 kwargs = vars(kwargs) 150 parser._defaults = old_defaults 151 152 if unknown: 153 if subcommand: 154 name = "{} {}".format(name, subcommand) 155 parser.error( 156 "unrecognized arguments for {}: {}".format( 157 name, ", ".join(["'{}'".format(arg) for arg in unknown]) 158 ) 159 ) 160 161 return self._run_command_handler(handler, context, **kwargs) 162 163 164Registrar = MachRegistrar() 165