From 79b3316325a5280b96854083f74d6d8e1760cb4c Mon Sep 17 00:00:00 2001 From: ahochheiden Date: Wed, 19 Jul 2023 05:09:01 +0000 Subject: [PATCH] Bug 1695312 - Activate the virtualenv associated with a mach command much earlier r=firefox-build-system-reviewers,glandium This activated virtualenv for a command is managed `CommandSiteManager` and it is passed down to where it was activated before to prevent a second, redundant, activation. Differential Revision: https://phabricator.services.mozilla.com/D180499 --- build/mach_initialize.py | 158 +++++++++++++++++++- mach | 10 +- python/mach/mach/main.py | 180 ++++++++++++----------- python/mach/mach/registrar.py | 13 +- python/mach/mach/test/test_dispatcher.py | 4 +- python/mozbuild/mozbuild/base.py | 4 + 6 files changed, 272 insertions(+), 97 deletions(-) diff --git a/build/mach_initialize.py b/build/mach_initialize.py index 3db4ba6d34d0..ce35cb801d77 100644 --- a/build/mach_initialize.py +++ b/build/mach_initialize.py @@ -137,7 +137,7 @@ def check_for_spaces(topsrcdir): ) -def initialize(topsrcdir): +def initialize(topsrcdir, args=()): # This directory was deleted in bug 1666345, but there may be some ignored # files here. We can safely just delete it for the user so they don't have # to clean the repo themselves. @@ -157,7 +157,7 @@ def initialize(topsrcdir): ) ] - from mach.util import get_state_dir, setenv + from mach.util import get_state_dir, get_virtualenv_base_dir, setenv state_dir = _create_state_dir() @@ -171,7 +171,7 @@ def initialize(topsrcdir): import mach.base import mach.main - from mach.main import MachCommandReference + from mach.main import MachCommandReference, get_argument_parser # Centralized registry of available mach commands MACH_COMMANDS = { @@ -407,6 +407,156 @@ def initialize(topsrcdir): "xpcshell-test": MachCommandReference("testing/xpcshell/mach_commands.py"), } + import argparse + import ast + + class DecoratorVisitor(ast.NodeVisitor): + def __init__(self): + self.results = {} + + def visit_FunctionDef(self, node): + # We only care about `Command` and `SubCommand` decorators, since + # they are the only ones that can specify virtualenv_name + decorators = [ + decorator + for decorator in node.decorator_list + if isinstance(decorator, ast.Call) + and isinstance(decorator.func, ast.Name) + and decorator.func.id in ["SubCommand", "Command"] + ] + + relevant_kwargs = ["command", "subcommand", "virtualenv_name"] + + for decorator in decorators: + kwarg_dict = {} + + for name, arg in zip(["command", "subcommand"], decorator.args): + kwarg_dict[name] = arg.s + + for keyword in decorator.keywords: + if keyword.arg not in relevant_kwargs: + # We only care about these 3 kwargs, so we can safely skip the rest + continue + + kwarg_dict[keyword.arg] = getattr(keyword.value, "s", "") + + command = kwarg_dict.pop("command") + self.results.setdefault(command, {}) + + sub_command = kwarg_dict.pop("subcommand", None) + virtualenv_name = kwarg_dict.pop("virtualenv_name", None) + + if sub_command: + self.results[command].setdefault("subcommands", {}) + sub_command_dict = self.results[command]["subcommands"].setdefault( + sub_command, {} + ) + + if virtualenv_name: + sub_command_dict["virtualenv_name"] = virtualenv_name + elif virtualenv_name: + # If there is no `subcommand` we are in the `@Command` + # decorator, and need to store the virtualenv_name for + # the 'command'. + self.results[command]["virtualenv_name"] = virtualenv_name + + self.generic_visit(node) + + def command_virtualenv_info_for_module(file_path): + command_module_path = Path(topsrcdir) / file_path + with command_module_path.open("r") as file: + content = file.read() + + tree = ast.parse(content) + visitor = DecoratorVisitor() + visitor.visit(tree) + + return visitor.results + + class DetermineCommandVenvAction(argparse.Action): + def __init__( + self, + option_strings, + dest, + required=True, + default=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. + # + # In a similar vein, default is passed in but is not needed, so we drop + # it. + argparse.Action.__init__( + self, + option_strings, + dest, + required=required, + help=argparse.SUPPRESS, + nargs=argparse.REMAINDER, + ) + + def __call__(self, parser, namespace, values, option_string=None): + if len(values) == 0: + return + + command = values[0] + setattr(namespace, "command_name", command) + + site = "common" + + if len(values) > 1: + potential_sub_command_name = values[1] + else: + potential_sub_command_name = None + + module_path = MACH_COMMANDS.get(command).module + + module_dict = command_virtualenv_info_for_module(module_path) + command_dict = module_dict.get(command, {}) + + if not command_dict: + return + + if ( + potential_sub_command_name + and not potential_sub_command_name.startswith("-") + ): + all_sub_commands_dict = command_dict.get("subcommands", {}) + + if all_sub_commands_dict: + sub_command_dict = all_sub_commands_dict.get( + potential_sub_command_name, {} + ) + + if sub_command_dict: + site = sub_command_dict.get("virtualenv_name", "common") + else: + site = command_dict.get("virtualenv_name", "common") + + setattr(namespace, "site_name", site) + + parser = get_argument_parser(action=DetermineCommandVenvAction) + namespace = parser.parse_args() + + command_name = getattr(namespace, "command_name", None) + site_name = getattr(namespace, "site_name", "common") + command_site_manager = None + + # the 'clobber' command needs to run in the 'mach' venv, so we + # don't want to activate any other virtualenv for it. + if command_name != "clobber": + from mach.site import CommandSiteManager + + command_site_manager = CommandSiteManager.from_environment( + topsrcdir, + lambda: os.path.normpath(get_state_dir(True, topsrcdir=topsrcdir)), + site_name, + get_virtualenv_base_dir(topsrcdir), + ) + + command_site_manager.activate() + # Set a reasonable limit to the number of open files. # # Some linux systems set `ulimit -n` to a very high number, which works @@ -530,7 +680,7 @@ def initialize(topsrcdir): if "MACH_MAIN_PID" not in os.environ: setenv("MACH_MAIN_PID", str(os.getpid())) - driver = mach.main.Mach(os.getcwd()) + driver = mach.main.Mach(os.getcwd(), command_site_manager) driver.populate_context_handler = populate_context if not driver.settings_paths: diff --git a/mach b/mach index cc5e25b7aa36..a62db70b13a5 100755 --- a/mach +++ b/mach @@ -13,17 +13,17 @@ MIN_PYTHON_VERSION = (3, 7) MAX_PYTHON_VERSION_TO_CONSIDER = (3, 11) -def load_mach(dir_path, mach_path): +def load_mach(dir_path, mach_path, args): # Defer import of "importlib.util" until after Python version check has happened # so that Python 2 usages fail gracefully. import importlib.util spec = importlib.util.spec_from_file_location('mach_initialize', mach_path) mach_initialize = importlib.util.module_from_spec(spec) spec.loader.exec_module(mach_initialize) - return mach_initialize.initialize(dir_path) + return mach_initialize.initialize(dir_path, args) -def check_and_get_mach(dir_path): +def check_and_get_mach(dir_path, args): initialize_paths = ( # Run Thunderbird's mach_initialize.py if it exists 'comm/build/mach_initialize.py', @@ -34,7 +34,7 @@ def check_and_get_mach(dir_path): for initialize_path in initialize_paths: mach_path = os.path.join(dir_path, initialize_path) if os.path.isfile(mach_path): - return load_mach(dir_path, mach_path) + return load_mach(dir_path, mach_path, args) return None @@ -114,7 +114,7 @@ def main(args): # https://github.com/python/cpython/pull/9516 os.environ.pop("__PYVENV_LAUNCHER__", None) - mach = check_and_get_mach(os.path.dirname(os.path.realpath(__file__))) + mach = check_and_get_mach(os.path.dirname(os.path.realpath(__file__)), args) if not mach: print('Could not run mach: No mach source directory found.') sys.exit(1) diff --git a/python/mach/mach/main.py b/python/mach/mach/main.py index 79fcb8e8ad17..f1c176fbe813 100644 --- a/python/mach/mach/main.py +++ b/python/mach/mach/main.py @@ -16,7 +16,9 @@ import traceback import uuid from collections.abc import Iterable from pathlib import Path -from typing import Dict, List, Union +from typing import Dict, List, Optional, Union + +from mach.site import CommandSiteManager from .base import ( CommandContext, @@ -226,7 +228,9 @@ To see more help for a specific command, run: %(prog)s help """ - def __init__(self, cwd: str): + def __init__( + self, cwd: str, command_site_manager: Optional[CommandSiteManager] = None + ): assert Path(cwd).is_dir() self.cwd = cwd @@ -234,6 +238,7 @@ To see more help for a specific command, run: self.logger = logging.getLogger(__name__) self.settings = ConfigSettings() self.settings_paths = [] + self.command_site_manager = command_site_manager if "MACHRC" in os.environ: self.settings_paths.append(os.environ["MACHRC"]) @@ -443,7 +448,7 @@ To see more help for a specific command, run: if self.populate_context_handler: context = ContextWrapper(context, self.populate_context_handler) - parser = self.get_argument_parser(context) + parser = get_argument_parser(context) context.global_parser = parser if not len(argv): @@ -519,6 +524,7 @@ To see more help for a specific command, run: return Registrar._run_command_handler( handler, context, + self.command_site_manager, debug_command=args.debug_command, profile_command=args.profile_command, **vars(args.command_args), @@ -641,95 +647,99 @@ To see more help for a specific command, run: self.settings.load_files(list(files)) - def get_argument_parser(self, context): - """Returns an argument parser for the command-line interface.""" - parser = ArgumentParser( - add_help=False, - usage="%(prog)s [global arguments] " "command [command arguments]", - ) +def get_argument_parser(context=None, action=CommandAction): + """Returns an argument parser for the command-line interface.""" - # WARNING!!! If you add a global argument here, also add it to the - # global argument handling in the top-level `mach` script. - # Order is important here as it dictates the order the auto-generated - # help messages are printed. - global_group = parser.add_argument_group("Global Arguments") + parser = ArgumentParser( + add_help=False, + usage="%(prog)s [global arguments] " "command [command arguments]", + ) - global_group.add_argument( - "-v", - "--verbose", - dest="verbose", - action="store_true", - default=False, - help="Print verbose output.", - ) - global_group.add_argument( - "-l", - "--log-file", - dest="logfile", - metavar="FILENAME", - type=argparse.FileType("a"), - help="Filename to write log data to.", - ) - global_group.add_argument( - "--log-interval", - dest="log_interval", - action="store_true", - default=False, - help="Prefix log line with interval from last message rather " - "than relative time. Note that this is NOT execution time " - "if there are parallel operations.", - ) - global_group.add_argument( - "--no-interactive", - dest="is_interactive", - action="store_false", - help="Automatically selects the default option on any " - "interactive prompts. If the output is not a terminal, " - "then --no-interactive is assumed.", - ) - suppress_log_by_default = False - if "INSIDE_EMACS" in os.environ: - suppress_log_by_default = True - global_group.add_argument( - "--log-no-times", - dest="log_no_times", - action="store_true", - default=suppress_log_by_default, - help="Do not prefix log lines with times. By default, " - "mach will prefix each output line with the time since " - "command start.", - ) - global_group.add_argument( - "-h", - "--help", - dest="help", - action="store_true", - default=False, - help="Show this help message.", - ) - global_group.add_argument( - "--debug-command", - action="store_true", - help="Start a Python debugger when command is dispatched.", - ) - global_group.add_argument( - "--profile-command", - action="store_true", - help="Capture a Python profile of the mach process as command is dispatched.", - ) - global_group.add_argument( - "--settings", - dest="settings_file", - metavar="FILENAME", - default=None, - help="Path to settings file.", - ) + # WARNING!!! If you add a global argument here, also add it to the + # global argument handling in the top-level `mach` script. + # Order is important here as it dictates the order the auto-generated + # help messages are printed. + global_group = parser.add_argument_group("Global Arguments") + global_group.add_argument( + "-v", + "--verbose", + dest="verbose", + action="store_true", + default=False, + help="Print verbose output.", + ) + global_group.add_argument( + "-l", + "--log-file", + dest="logfile", + metavar="FILENAME", + type=argparse.FileType("a"), + help="Filename to write log data to.", + ) + global_group.add_argument( + "--log-interval", + dest="log_interval", + action="store_true", + default=False, + help="Prefix log line with interval from last message rather " + "than relative time. Note that this is NOT execution time " + "if there are parallel operations.", + ) + global_group.add_argument( + "--no-interactive", + dest="is_interactive", + action="store_false", + help="Automatically selects the default option on any " + "interactive prompts. If the output is not a terminal, " + "then --no-interactive is assumed.", + ) + suppress_log_by_default = False + if "INSIDE_EMACS" in os.environ: + suppress_log_by_default = True + global_group.add_argument( + "--log-no-times", + dest="log_no_times", + action="store_true", + default=suppress_log_by_default, + help="Do not prefix log lines with times. By default, " + "mach will prefix each output line with the time since " + "command start.", + ) + global_group.add_argument( + "-h", + "--help", + dest="help", + action="store_true", + default=False, + help="Show this help message.", + ) + global_group.add_argument( + "--debug-command", + action="store_true", + help="Start a Python debugger when command is dispatched.", + ) + global_group.add_argument( + "--profile-command", + action="store_true", + help="Capture a Python profile of the mach process as command is dispatched.", + ) + global_group.add_argument( + "--settings", + dest="settings_file", + metavar="FILENAME", + default=None, + help="Path to settings file.", + ) + + if context: # 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, context=context ) + else: + parser.add_argument("command", action=action) - return parser + return parser diff --git a/python/mach/mach/registrar.py b/python/mach/mach/registrar.py index 75481596f4b4..e7738af81c31 100644 --- a/python/mach/mach/registrar.py +++ b/python/mach/mach/registrar.py @@ -86,7 +86,13 @@ class MachRegistrar(object): return fail_conditions def _run_command_handler( - self, handler, context, debug_command=False, profile_command=False, **kwargs + self, + handler, + context, + command_site_manager=None, + debug_command=False, + profile_command=False, + **kwargs, ): instance = MachRegistrar._instance(handler, context, **kwargs) fail_conditions = MachRegistrar._fail_conditions(handler, instance) @@ -99,7 +105,10 @@ class MachRegistrar(object): self.command_depth += 1 fn = handler.func if handler.virtualenv_name: - instance.activate_virtualenv() + if command_site_manager: + instance.virtualenv_manager = command_site_manager + else: + instance.activate_virtualenv() profile = None if profile_command: diff --git a/python/mach/mach/test/test_dispatcher.py b/python/mach/mach/test/test_dispatcher.py index 85c2e9a847d8..b88b58ba745d 100644 --- a/python/mach/mach/test/test_dispatcher.py +++ b/python/mach/mach/test/test_dispatcher.py @@ -30,7 +30,9 @@ class TestDispatcher(unittest.TestCase): mach.settings.load_fps([config]) context = CommandContext(cwd="", settings=mach.settings) - return mach.get_argument_parser(context) + from mach.main import get_argument_parser + + return get_argument_parser(context) def test_command_aliases(self): config = """ diff --git a/python/mozbuild/mozbuild/base.py b/python/mozbuild/mozbuild/base.py index 8dc9f90dbf64..3f5a6242f838 100644 --- a/python/mozbuild/mozbuild/base.py +++ b/python/mozbuild/mozbuild/base.py @@ -268,6 +268,10 @@ class MozbuildObject(ProcessExecutionMixin): return self._virtualenv_manager + @virtualenv_manager.setter + def virtualenv_manager(self, command_site_manager): + self._virtualenv_manager = command_site_manager + @staticmethod @memoize def get_base_mozconfig_info(topsrcdir, path, env_mozconfig):