Source code
Revision control
Copy as Markdown
Other Tools
# 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,
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Optional
import tomlkit
from mozfile import which
from mozbuild.base import MozbuildObject
from mozbuild.lockfiles.site_dependency_extractor import SiteDependencyExtractor
class MissingMachSiteFileError(Exception):
"""Raised when the required mach.txt site file is missing."""
pass
class MissingUVError(Exception):
"""Raised when the required 'uv' executable is not on PATH."""
pass
class GeneratePythonLockfiles(MozbuildObject):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, virtualenv_name="uv", **kwargs)
self.keep_lockfiles = False
self.output_dir = None
self.sites_dir = Path(self.topsrcdir) / "python" / "sites"
self.command_site_names_to_lock = {}
self.mach_dependencies = []
self.return_code = 0
def setup(self, keep_lockfiles: bool, sites: Optional[list[str]] = None) -> None:
self.keep_lockfiles = keep_lockfiles
all_sites = {
p.stem
for p in self.sites_dir.iterdir()
if p.is_file() and p.suffix == ".txt"
}
if sites:
sites = set(sites)
else:
sites = all_sites
# We will always create a lockfile for the `mach` site, so let's
# make this set only for the command `sites`.
self.command_site_names_to_lock = sites
self.command_site_names_to_lock.discard("mach")
# It should be impossible to get here without a `mach.txt`
# file, but we'll do a sanity check anyway.
mach_file = self.sites_dir / "mach.txt"
if not mach_file.exists():
raise MissingMachSiteFileError(
f"Required site file 'mach.txt' not found in {self.sites_dir}"
)
# UV should be installed with the virtualenv, but we'll verify anyway.
if which("uv") is None:
raise MissingUVError("'uv' executable not found in PATH")
def cleanup(self):
if self.keep_lockfiles:
print(f"\nGenerated lockfiles are retained in {self.output_dir}")
else:
shutil.rmtree(self.output_dir)
def create_pyproject_toml_for_site(self, site_name: str):
site_dir = self.output_dir / site_name
site_dir.mkdir()
subprocess.check_call(
[
"uv",
"init",
f"--name={site_name}-lock",
"--bare",
f"--description=Used to generate the lockfile for the {site_name} site.",
"--no-progress",
],
cwd=site_dir,
)
return site_dir / "pyproject.toml"
def lock_site(self, site_name: str):
print(f"\nPreparing to lock the `{site_name}` site")
extractor = SiteDependencyExtractor(
site_name=site_name, sites_dir=self.sites_dir, topsrcdir=self.topsrcdir
)
requires_python, dependencies = extractor.parse()
# We always lock the mach site first and save its dependencies.
# For all other sites, append those saved dependencies since they inherit from mach.
if site_name == "mach":
self.mach_dependencies = dependencies
else:
dependencies += self.mach_dependencies
pyproject_toml_path = self.create_pyproject_toml_for_site(site_name=site_name)
toml_doc = tomlkit.parse(pyproject_toml_path.read_text(encoding="utf-8"))
project = toml_doc.setdefault("project", tomlkit.table())
project["requires-python"] = requires_python
deps = tomlkit.array()
deps.multiline(True)
deps.indent(4)
for dep in dependencies:
deps.append(dep.name + dep.version)
project["dependencies"] = deps
pyproject_toml_path.write_text(tomlkit.dumps(toml_doc), encoding="utf-8")
print(f"Attempting to create a lockfile for the `{site_name}` site")
result = subprocess.run(
["uv", "lock", "--no-progress"], cwd=pyproject_toml_path.parent, check=False
)
if result.returncode:
print(f"Failed to create lockfile for `{site_name}` site...")
if not self.return_code:
self.return_code = result.returncode
else:
print(f"Successfully created lockfile for {site_name}")
def run(self, keep_lockfiles: bool, sites: Optional[list[str]] = None) -> int:
try:
self.output_dir = Path(tempfile.mkdtemp(prefix="mach_python_lockfiles_"))
self.setup(keep_lockfiles, sites)
self.lock_site(site_name="mach")
for site_name in self.command_site_names_to_lock:
self.lock_site(site_name=site_name)
finally:
self.cleanup()
return self.return_code