1 files changed, 120 insertions, 0 deletions
diff --git a/comm/python/mutlh/mutlh/site.py b/comm/python/mutlh/mutlh/site.py
new file mode 100644
index 0000000000..ab93e874e9
--- /dev/null
+++ b/comm/python/mutlh/mutlh/site.py
@@ -0,0 +1,120 @@
+#  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 https://mozilla.org/MPL/2.0/.
+
+import functools
+import os
+from typing import Callable, Optional
+
+from mach.requirements import MachEnvRequirements, UnexpectedFlexibleRequirementException
+from mach.site import (
+    PIP_NETWORK_INSTALL_RESTRICTED_VIRTUALENVS,
+    CommandSiteManager,
+    MozSiteMetadata,
+    SitePackagesSource,
+    _mach_virtualenv_root,
+)
+
+
+class SiteNotFoundException(Exception):
+    def __init__(self, site_name, manifest_paths):
+        self.site_name = site_name
+        self.manifest_paths = manifest_paths
+        self.args = (site_name, manifest_paths)
+
+
+@functools.lru_cache(maxsize=None)
+def find_manifest(topsrcdir, site_name):
+    manifest_paths = (
+        os.path.join(topsrcdir, "comm", "python", "sites", f"{site_name}.txt"),
+        os.path.join(topsrcdir, "python", "sites", f"{site_name}.txt"),
+    )
+
+    for check_path in manifest_paths:
+        if os.path.exists(check_path):
+            return check_path
+
+    raise SiteNotFoundException(site_name, manifest_paths)
+
+
+@functools.lru_cache(maxsize=None)
+def resolve_requirements(topsrcdir, site_name):
+    try:
+        manifest_path = find_manifest(topsrcdir, site_name)
+    except SiteNotFoundException as e:
+        raise Exception(
+            f'The current command is using the "{e.site_name}" '
+            "site. However, that site is missing its associated "
+            f"requirements definition file in one of the supported "
+            f"paths: {e.manifest_paths}."
+        )
+    is_thunderbird = True
+    try:
+        return MachEnvRequirements.from_requirements_definition(
+            topsrcdir,
+            is_thunderbird,
+            site_name not in PIP_NETWORK_INSTALL_RESTRICTED_VIRTUALENVS,
+            manifest_path,
+        )
+    except UnexpectedFlexibleRequirementException as e:
+        raise Exception(
+            f'The "{site_name}" site does not have all pypi packages pinned '
+            f'in the format "package==version" (found "{e.raw_requirement}").\n'
+            f"Only the {PIP_NETWORK_INSTALL_RESTRICTED_VIRTUALENVS} sites are "
+            "allowed to have unpinned packages."
+        )
+
+
+class MutlhCommandSiteManager(CommandSiteManager):
+    @classmethod
+    def from_environment(
+        cls,
+        topsrcdir: str,
+        get_state_dir: Callable[[], Optional[str]],
+        site_name: str,
+        command_virtualenvs_dir: str,
+    ):
+        """
+        Args:
+            topsrcdir: The path to the Firefox repo
+            get_state_dir: A function that resolves the path to the checkout-scoped
+                state_dir, generally ~/.mozbuild/srcdirs/<checkout-based-dir>/
+            site_name: The name of this site, such as "build"
+            command_virtualenvs_dir: The location under which this site's virtualenv
+            should be created
+        """
+        active_metadata = MozSiteMetadata.from_runtime()
+        assert (
+            active_metadata
+        ), "A Mach-managed site must be active before doing work with command sites"
+
+        mach_site_packages_source = active_metadata.mach_site_packages_source
+        pip_restricted_site = site_name in PIP_NETWORK_INSTALL_RESTRICTED_VIRTUALENVS
+        if not pip_restricted_site and mach_site_packages_source == SitePackagesSource.SYSTEM:
+            # Sites that aren't pip-network-install-restricted are likely going to be
+            # incompatible with the system. Besides, this use case shouldn't exist, since
+            # using the system packages is supposed to only be needed to lower risk of
+            # important processes like building Firefox.
+            raise Exception(
+                'Cannot use MACH_BUILD_PYTHON_NATIVE_PACKAGE_SOURCE="system" for any '
+                f"sites other than {PIP_NETWORK_INSTALL_RESTRICTED_VIRTUALENVS}. The "
+                f'current attempted site is "{site_name}".'
+            )
+
+        mach_virtualenv_root = (
+            _mach_virtualenv_root(get_state_dir())
+            if mach_site_packages_source == SitePackagesSource.VENV
+            else None
+        )
+        populate_virtualenv = (
+            mach_site_packages_source == SitePackagesSource.VENV or not pip_restricted_site
+        )
+        return cls(
+            topsrcdir,
+            mach_virtualenv_root,
+            os.path.join(command_virtualenvs_dir, site_name),
+            site_name,
+            active_metadata,
+            populate_virtualenv,
+            resolve_requirements(topsrcdir, site_name),
+        )