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

@@ -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:

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)

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 <command>
 """
 
-    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,7 +647,8 @@ To see more help for a specific command, run:
 
         self.settings.load_files(list(files))
 
-    def get_argument_parser(self, context):
+
+def get_argument_parser(context=None, action=CommandAction):
     """Returns an argument parser for the command-line interface."""
 
     parser = ArgumentParser(
@@ -726,10 +733,13 @@ To see more help for a specific command, run:
         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

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,6 +105,9 @@ class MachRegistrar(object):
         self.command_depth += 1
         fn = handler.func
         if handler.virtualenv_name:
+            if command_site_manager:
+                instance.virtualenv_manager = command_site_manager
+            else:
                 instance.activate_virtualenv()
 
         profile = None

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 = """

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):