Bug 1717051: Reuse "requirements" logic in mach_initialize r=ahal

Rather than re-implementing it as `search_path()`, use the existing
`MachEnvRequirements` tool to parse `mach_virtualenv_requirements.txt`

Differential Revision: https://phabricator.services.mozilla.com/D126280

build/mach_initialize.py

@@ -141,35 +141,6 @@ CATEGORIES = {
     },
 }
 
-
-def search_path(mozilla_dir, packages_txt):
-    with open(os.path.join(mozilla_dir, packages_txt)) as f:
-        packages = [
-            line.strip().split(":", maxsplit=1)
-            for line in f
-            if not line.lstrip().startswith("#")
-        ]
-
-    def handle_package(action, package):
-        if action == "packages.txt":
-            for p in search_path(mozilla_dir, package):
-                yield os.path.join(mozilla_dir, p)
-
-        if action == "pth":
-            yield os.path.join(mozilla_dir, package)
-
-    for current_action, current_package in packages:
-        for path in handle_package(current_action, current_package):
-            yield path
-
-
-def mach_sys_path(mozilla_dir):
-    return [
-        os.path.join(mozilla_dir, path)
-        for path in search_path(mozilla_dir, "build/mach_virtualenv_packages.txt")
-    ]
-
-
 INSTALL_PYTHON_GUIDANCE_LINUX = """
 See https://firefox-source-docs.mozilla.org/setup/linux_build.html#installingpython
 for guidance on how to install Python on your system.
@@ -196,6 +167,28 @@ install a recent enough Python 3.
 """.strip()
 
 
+def _activate_python_environment(topsrcdir):
+    # We need the "mach" module to access the logic to parse virtualenv
+    # requirements.
+    sys.path.insert(0, os.path.join(topsrcdir, "python", "mach"))
+
+    from mach.requirements import MachEnvRequirements
+
+    thunderbird_dir = os.path.join(topsrcdir, "comm")
+    is_thunderbird = os.path.exists(thunderbird_dir) and bool(
+        os.listdir(thunderbird_dir)
+    )
+
+    requirements = MachEnvRequirements.from_requirements_definition(
+        topsrcdir,
+        is_thunderbird,
+        os.path.join(topsrcdir, "build", "mach_virtualenv_packages.txt"),
+    )
+    sys.path[0:0] = [
+        os.path.join(topsrcdir, pth.path) for pth in requirements.pth_requirements
+    ]
+
+
 def initialize(topsrcdir):
     # Ensure we are running Python 3.6+. We run this check as soon as
     # possible to avoid a cryptic import/usage error.
@@ -226,8 +219,8 @@ def initialize(topsrcdir):
         sys.path = [path for path in sys.path if path not in site_paths]
 
     state_dir = _create_state_dir()
+    _activate_python_environment(topsrcdir)
 
-    sys.path[0:0] = mach_sys_path(topsrcdir)
     import mach.base
     import mach.main
     from mach.util import setenv

configure.py

@@ -21,6 +21,7 @@ except ImportError:
 
 
 base_dir = os.path.abspath(os.path.dirname(__file__))
+sys.path.insert(0, os.path.join(base_dir, "python", "mach"))
 sys.path.insert(0, os.path.join(base_dir, "python", "mozboot"))
 sys.path.insert(0, os.path.join(base_dir, "python", "mozbuild"))
 sys.path.insert(0, os.path.join(base_dir, "third_party", "python", "six"))

python/mach_commands.py

@@ -22,6 +22,7 @@ from manifestparser import filters as mpf
 
 
 from mach.decorators import CommandArgument, Command
+from mach.requirements import MachEnvRequirements
 from mach.util import UserError
 
 here = os.path.abspath(os.path.dirname(__file__))
@@ -68,11 +69,18 @@ def python(
         raise UserError("Cannot pass both --requirements and --no-virtualenv.")
 
     if no_virtualenv:
-        from mach_initialize import mach_sys_path
-
         python_path = sys.executable
+        requirements = MachEnvRequirements.from_requirements_definition(
+            command_context.topsrcdir,
+            False,
+            os.path.join(
+                command_context.topsrcdir, "build", "mach_virtualenv_packages.txt"
+            ),
+        )
+
         append_env["PYTHONPATH"] = os.pathsep.join(
-            mach_sys_path(command_context.topsrcdir)
+            os.path.join(command_context.topsrcdir, pth.path)
+            for pth in requirements.pth_requirements
         )
     else:
         command_context.virtualenv_manager.ensure()

python/mozbuild/mozbuild/test/test_vendor.py

@@ -37,6 +37,13 @@ def test_up_to_date_vendor():
         work_vendored = os.path.join(work_dir, "third_party", "python")
         shutil.copytree(existing_vendored, work_vendored)
 
+        # Copy "mach" module so that `VirtualenvManager` can populate itself.
+        # This is needed because "topsrcdir" is used in this test both for determining
+        # import paths and for acting as a "work dir".
+        existing_mach = os.path.join(topsrcdir, "python", "mach")
+        work_mach = os.path.join(work_dir, "python", "mach")
+        shutil.copytree(existing_mach, work_mach)
+
         # Run the vendoring process
         vendor = VendorPython(
             work_dir, None, Mock(), topobjdir=os.path.join(work_dir, "obj")

python/mozbuild/mozbuild/virtualenv.py

@@ -309,15 +309,7 @@ class VirtualenvManager(VirtualenvHelper):
         return self.virtualenv_root
 
     def _requirements(self):
-        try:
-            # When `virtualenv.py` is invoked from an existing Mach process,
-            # import MachEnvRequirements in the expected way.
-            from mozbuild.requirements import MachEnvRequirements
-        except ImportError:
-            # When `virtualenv.py` is invoked standalone, import
-            # MachEnvRequirements from the adjacent "standalone"
-            # requirements module.
-            from requirements import MachEnvRequirements
+        from mach.requirements import MachEnvRequirements
 
         if not os.path.exists(self._manifest_path):
             raise Exception(
@@ -603,6 +595,9 @@ if __name__ == "__main__":
         populate = False
         opts = parser.parse_args(sys.argv[1:])
 
+    # We want to be able to import the "mach.requirements" module.
+    sys.path.append(os.path.join(opts.topsrcdir, "python", "mach"))
+
     manager = VirtualenvManager(
         opts.topsrcdir,
         opts.virtualenvs_dir,