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