Bug 1723031: Allow flexible dependency-specification in the Mach venv r=ahal

There's some trade-offs in play here: the major issue is that we can't
pin `psutil`'s because it's pre-installed on some CI workers with a
different version (5.4.2).

Additionally, we can't "just" bump that version, because CI workers jump
between different revisions to do work, so a specific pinned version
won't work when we try to update such a package.

One option is that we could avoid validating package versions in CI, but
then that will cause surprises (heck, I didn't know we were still using
`psutil==5.4.2` instead of `5.8.0` until now). By doing validation, we
make it more explicit and avoid accidentally depending on behaviour of
too new of such a package.

However, in most cases, we manage the installation environment and can
pin dependencies. So, I've made the top-level Mach virtualenv the _only_
one that is able to use requirement specifiers other than "==".

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

build/mach_initialize.py

@@ -169,8 +169,16 @@ install a recent enough Python 3.
 
 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"))
+    # requirements. Since that depends on "packaging" (and, transitively,
+    # "pyparsing"), we add those to the path too.
+    sys.path[0:0] = [
+        os.path.join(topsrcdir, module)
+        for module in (
+            os.path.join("python", "mach"),
+            os.path.join("third_party", "python", "packaging"),
+            os.path.join("third_party", "python", "pyparsing"),
+        )
+    ]
 
     from mach.requirements import MachEnvRequirements
 
@@ -182,6 +190,7 @@ def _activate_python_environment(topsrcdir):
     requirements = MachEnvRequirements.from_requirements_definition(
         topsrcdir,
         is_thunderbird,
+        True,
         os.path.join(topsrcdir, "build", "mach_virtualenv_packages.txt"),
     )
     sys.path[0:0] = [

build/mach_virtualenv_packages.txt

@@ -3,5 +3,7 @@ packages.txt:build/common_virtualenv_packages.txt
 # and it has to be built from source.
 pypi-optional:glean-sdk==40.0.0:telemetry will not be collected
 # Mach gracefully handles the case where `psutil` is unavailable.
-pypi-optional:psutil==5.8.0:telemetry will be missing some data
+# We aren't (yet) able to pin packages in automation, so we have to
+# support down to the oldest locally-installed version (5.4.2).
+pypi-optional:psutil>=5.4.2,<=5.8.0:telemetry will be missing some data
 pypi:zstandard==0.15.2

configure.py

@@ -24,6 +24,8 @@ 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", "packaging"))
+sys.path.insert(0, os.path.join(base_dir, "third_party", "python", "pyparsing"))
 sys.path.insert(0, os.path.join(base_dir, "third_party", "python", "six"))
 from mozbuild.configure import (
     ConfigureSandbox,

python/mach/mach/requirements.py

@@ -5,6 +5,8 @@
 import os
 from pathlib import Path
 
+from packaging.requirements import Requirement
+
 
 THUNDERBIRD_PYPI_ERROR = """
 Thunderbird requirements definitions cannot include PyPI packages.
@@ -17,18 +19,14 @@ class PthSpecifier:
 
 
 class PypiSpecifier:
-    def __init__(self, package_name, version, full_specifier):
-        self.package_name = package_name
-        self.version = version
-        self.full_specifier = full_specifier
+    def __init__(self, requirement):
+        self.requirement = requirement
 
 
-class PypiOptionalSpecifier:
-    def __init__(self, repercussion, package_name, version, full_specifier):
+class PypiOptionalSpecifier(PypiSpecifier):
+    def __init__(self, repercussion, requirement):
+        super().__init__(requirement)
         self.repercussion = repercussion
-        self.package_name = package_name
-        self.version = version
-        self.full_specifier = full_specifier
 
 
 class MachEnvRequirements:
@@ -65,17 +63,29 @@ class MachEnvRequirements:
 
     @classmethod
     def from_requirements_definition(
-        cls, topsrcdir, is_thunderbird, requirements_definition
+        cls,
+        topsrcdir,
+        is_thunderbird,
+        is_mach_or_build_virtualenv,
+        requirements_definition,
     ):
         requirements = cls()
         _parse_mach_env_requirements(
-            requirements, requirements_definition, topsrcdir, is_thunderbird
+            requirements,
+            requirements_definition,
+            topsrcdir,
+            is_thunderbird,
+            is_mach_or_build_virtualenv,
         )
         return requirements
 
 
 def _parse_mach_env_requirements(
-    requirements_output, root_requirements_path, topsrcdir, is_thunderbird
+    requirements_output,
+    root_requirements_path,
+    topsrcdir,
+    is_thunderbird,
+    is_mach_or_build_virtualenv,
 ):
     topsrcdir = Path(topsrcdir)
 
@@ -119,9 +129,10 @@ def _parse_mach_env_requirements(
             if is_thunderbird_packages_txt:
                 raise Exception(THUNDERBIRD_PYPI_ERROR)
 
-            package_name, version = _parse_package_specifier(params)
             requirements_output.pypi_requirements.append(
-                PypiSpecifier(package_name, version, params)
+                PypiSpecifier(
+                    _parse_package_specifier(params, is_mach_or_build_virtualenv)
+                )
             )
         elif action == "pypi-optional":
             if is_thunderbird_packages_txt:
@@ -133,10 +144,14 @@ def _parse_mach_env_requirements(
                     'description in the format "package:fallback explanation", '
                     'found "{}"'.format(params)
                 )
-            package, repercussion = params.split(":")
-            package_name, version = _parse_package_specifier(package)
+            raw_requirement, repercussion = params.split(":")
             requirements_output.pypi_optional_requirements.append(
-                PypiOptionalSpecifier(repercussion, package_name, version, package)
+                PypiOptionalSpecifier(
+                    repercussion,
+                    _parse_package_specifier(
+                        raw_requirement, is_mach_or_build_virtualenv
+                    ),
+                )
             )
         elif action == "thunderbird-packages.txt":
             if is_thunderbird:
@@ -164,10 +179,14 @@ def _parse_mach_env_requirements(
     _parse_requirements_definition_file(root_requirements_path, False)
 
 
-def _parse_package_specifier(specifier):
-    if len(specifier.split("==")) != 2:
+def _parse_package_specifier(raw_requirement, is_mach_or_build_virtualenv):
+    requirement = Requirement(raw_requirement)
+
+    if not is_mach_or_build_virtualenv and [
+        s for s in requirement.specifier if s.operator != "=="
+    ]:
         raise Exception(
-            "Expected pypi package version to be pinned in the "
-            'format "package==version", found "{}"'.format(specifier)
+            'All virtualenvs except for "mach" and "build" must pin pypi package '
+            f'versions in the format "package==version", found "{raw_requirement}"'
         )
-    return specifier.split("==")
+    return requirement

python/mach/mach/test/test_virtualenv_compatibility.py

@@ -26,18 +26,18 @@ def _resolve_command_virtualenv_names():
     return virtualenv_names
 
 
-def _requirement_definition_to_pip_format(virtualenv_name, cache, is_mach_env):
+def _requirement_definition_to_pip_format(virtualenv_name, cache, is_mach_or_build_env):
     """Convert from parsed requirements object to pip-consumable format"""
     path = Path(topsrcdir) / "build" / f"{virtualenv_name}_virtualenv_packages.txt"
     requirements = MachEnvRequirements.from_requirements_definition(
-        topsrcdir, False, path
+        topsrcdir, False, is_mach_or_build_env, path
     )
 
     lines = []
     for pypi in (
         requirements.pypi_requirements + requirements.pypi_optional_requirements
     ):
-        lines.append(pypi.full_specifier)
+        lines.append(str(pypi.requirement))
 
     for vendored in requirements.vendored_requirements:
         lines.append(cache.package_for_vendor_dir(Path(vendored.path)))
@@ -111,7 +111,9 @@ def test_virtualenvs_compatible(tmpdir):
 
     for name in command_virtualenv_names:
         print(f'Checking compatibility of "{name}" virtualenv')
-        command_requirements = _requirement_definition_to_pip_format(name, cache, False)
+        command_requirements = _requirement_definition_to_pip_format(
+            name, cache, name == "build"
+        )
         with open(work_dir / "requirements.txt", "w") as requirements_txt:
             requirements_txt.write(mach_requirements)
             requirements_txt.write("\n")

python/mach_commands.py

@@ -73,6 +73,7 @@ def python(
         requirements = MachEnvRequirements.from_requirements_definition(
             command_context.topsrcdir,
             False,
+            True,
             os.path.join(
                 command_context.topsrcdir, "build", "mach_virtualenv_packages.txt"
             ),

python/mozbuild/mozbuild/virtualenv.py

@@ -233,18 +233,17 @@ class VirtualenvManager(VirtualenvHelper):
             installed_packages = {
                 package["name"]: package["version"] for package in installed_packages
             }
-            for requirement in env_requirements.pypi_requirements:
-                if (
-                    installed_packages.get(requirement.package_name, None)
-                    != requirement.version
+            for pkg in env_requirements.pypi_requirements:
+                if not pkg.requirement.specifier.contains(
+                    installed_packages.get(pkg.requirement.name, None)
                 ):
                     return False
 
-            for requirement in env_requirements.pypi_optional_requirements:
-                installed_version = installed_packages.get(
-                    requirement.package_name, None
-                )
-                if installed_version and installed_version != requirement.version:
+            for pkg in env_requirements.pypi_optional_requirements:
+                installed_version = installed_packages.get(pkg.requirement.name, None)
+                if installed_version and not pkg.requirement.specifier.contains(
+                    installed_packages.get(pkg.requirement.name, None)
+                ):
                     return False
 
         return True
@@ -326,7 +325,10 @@ class VirtualenvManager(VirtualenvHelper):
             os.listdir(thunderbird_dir)
         )
         return MachEnvRequirements.from_requirements_definition(
-            self.topsrcdir, is_thunderbird, self._manifest_path
+            self.topsrcdir,
+            is_thunderbird,
+            self._virtualenv_name in ("mach", "build"),
+            self._manifest_path,
         )
 
     def populate(self):
@@ -370,14 +372,14 @@ class VirtualenvManager(VirtualenvHelper):
                         f.write("{}\n".format(os.path.relpath(path, python_lib)))
 
             for pypi_requirement in env_requirements.pypi_requirements:
-                self.install_pip_package(pypi_requirement.full_specifier)
+                self.install_pip_package(str(pypi_requirement.requirement))
 
             for requirement in env_requirements.pypi_optional_requirements:
                 try:
-                    self.install_pip_package(requirement.full_specifier)
+                    self.install_pip_package(str(requirement.requirement))
                 except subprocess.CalledProcessError:
                     print(
-                        f"Could not install {requirement.package_name}, so "
+                        f"Could not install {requirement.requirement.name}, so "
                         f"{requirement.repercussion}. Continuing."
                     )
 
@@ -603,6 +605,9 @@ if __name__ == "__main__":
 
     # We want to be able to import the "mach.requirements" module.
     sys.path.append(os.path.join(opts.topsrcdir, "python", "mach"))
+    # Virtualenv logic needs access to the vendored "packaging" library.
+    sys.path.append(os.path.join(opts.topsrcdir, "third_party", "python", "pyparsing"))
+    sys.path.append(os.path.join(opts.topsrcdir, "third_party", "python", "packaging"))
 
     manager = VirtualenvManager(
         opts.topsrcdir,