Lock command PoC

src/pip_deepfreeze/__main__.py

@@ -6,6 +6,7 @@
 import typer
 from packaging.utils import canonicalize_name
 
+from .lock import lock as lock_operation
 from .pyproject_toml import load_pyproject_toml
 from .sanity import check_env
 from .sync import sync as sync_operation
@@ -88,6 +89,59 @@ def sync(
     )
 
 
+@app.command()
+def lock(
+    ctx: typer.Context,
+    to_upgrade: str = typer.Option(
+        None,
+        "--update",
+        "-u",
+        metavar="DEP1,DEP2,...",
+        help=(
+            "Make sure selected dependencies are upgraded (or downgraded) to "
+            "the latest allowed version. If DEP is not part of your application "
+            "dependencies anymore, this option has no effect."
+        ),
+    ),
+    upgrade_all: bool = typer.Option(
+        False,
+        "--update-all",
+        help=(
+            "Upgrade (or downgrade) all dependencies of your application to "
+            "the latest allowed version."
+        ),
+        show_default=False,
+    ),
+    extras: str = typer.Option(
+        None,
+        "--extras",
+        "-x",
+        metavar="EXTRAS",
+        help=(
+            "Comma separated list of extras "
+            "to install and freeze to requirements-{EXTRA}.txt."
+        ),
+    ),
+    post_lock_commands: List[str] = typer.Option(
+        [],
+        "--post-lock-command",
+        help=(
+            "Command to run after the lock operation is complete. "
+            "Can be specified multiple times."
+        ),
+    ),
+) -> None:
+    """Generate/update lock files without modifying the environment."""
+    lock_operation(
+        ctx.obj.python,
+        upgrade_all,
+        comma_split(to_upgrade),
+        extras=[canonicalize_name(extra) for extra in comma_split(extras)],
+        project_root=ctx.obj.project_root,
+        post_lock_commands=post_lock_commands,
+    )
+
+
 @app.command()
 def tree(
     ctx: typer.Context,

src/pip_deepfreeze/lock.py

@@ -0,0 +1,101 @@
+import subprocess
+from pathlib import Path
+from typing import Iterator, List, Optional, Sequence
+
+from packaging.utils import NormalizedName
+
+from .installed_dist import installed_distributions_as_pip_requirements
+from .pip import (
+    _dependencies_by_extra,  # TODO _dependencies_by_extra should be elsewhere
+    pip_dry_run_install_project,
+)
+from .project_name import get_project_name
+from .req_file_parser import OptionsLine, parse as parse_req_file
+from .req_merge import prepare_frozen_reqs_for_upgrade
+from .utils import (
+    HttpFetcher,
+    get_temp_path_in_dir,
+    log_info,
+    open_with_rollback,
+)
+
+
+def _make_requirements_path(project_root: Path, extra: Optional[str]) -> Path:
+    if extra:
+        return project_root / f"requirements-{extra}.txt"
+    else:
+        return project_root / "requirements.txt"
+
+
+def _make_requirements_paths(
+    project_root: Path, extras: Sequence[str]
+) -> Iterator[Path]:
+    yield _make_requirements_path(project_root, None)
+    for extra in extras:
+        yield _make_requirements_path(project_root, extra)
+
+
+def lock(
+    python: str,
+    upgrade_all: bool,
+    to_upgrade: List[str],
+    extras: List[NormalizedName],
+    project_root: Path,
+    post_lock_commands: Sequence[str] = (),
+) -> None:
+    project_name = get_project_name(python, project_root)
+    requirements_in = project_root / "requirements.txt.in"
+    # upgrade project and its dependencies, if needed
+    constraints_path = get_temp_path_in_dir(
+        dir=project_root, prefix="requirements.", suffix=".txt.df"
+    )
+    with constraints_path.open(mode="w", encoding="utf-8") as constraints:
+        for req_line in prepare_frozen_reqs_for_upgrade(
+            _make_requirements_paths(project_root, extras),
+            requirements_in,
+            upgrade_all,
+            to_upgrade,
+        ):
+            print(req_line, file=constraints)
+    installed_distributions = pip_dry_run_install_project(
+        python,
+        constraints_path,
+        project_root,
+        extras=extras,
+    )
+    # freeze dependencies
+    # TODO _dependencies_by_extra should be elsewhere
+    frozen_reqs_by_extra, unneeded_reqs = _dependencies_by_extra(
+        project_name,
+        extras,
+        installed_distributions,
+        installed_distributions_as_pip_requirements(installed_distributions),
+    )
+    assert not unneeded_reqs
+    for extra, frozen_reqs in frozen_reqs_by_extra.items():
+        requirements_frozen_path = _make_requirements_path(project_root, extra)
+        with open_with_rollback(requirements_frozen_path) as f:
+            print("# frozen requirements generated by pip-deepfreeze", file=f)
+            # output pip options in main requirements only
+            if not extra and requirements_in.exists():
+                # XXX can we avoid this second parse of requirements.txt.in?
+                for parsed_req_line in parse_req_file(
+                    str(requirements_in),
+                    reqs_only=False,
+                    recurse=True,
+                    strict=True,
+                    http_fetcher=HttpFetcher(),
+                ):
+                    if isinstance(parsed_req_line, OptionsLine):
+                        print(parsed_req_line.raw_line, file=f)
+            # output frozen dependencies of project
+            for req_line in frozen_reqs:
+                print(req_line, file=f)
+    # run post-lock commands
+    for command in post_lock_commands:
+        log_info(f"Running post-lock command: {command}")
+        result = subprocess.run(command, shell=True, check=False)
+        if result.returncode != 0:
+            raise SystemExit(
+                f"Post-lock command {command} failed with exit code {result.returncode}"
+            )

src/pip_deepfreeze/pip.py

@@ -1,5 +1,6 @@
 import json
 import shlex
+import tempfile
 from importlib.resources import path as resource_path
 from pathlib import Path
 from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, TypedDict, cast
@@ -11,6 +12,7 @@
     EnvInfoInstalledDistribution,
     InstalledDistributions,
     PipInspectInstalledDistribution,
+    PipInstallReportItemInstalledDistribution,
 )
 from .list_installed_depends import (
     list_installed_depends,
@@ -41,6 +43,64 @@ class PipInspectReport(TypedDict, total=False):
     environment: Dict[str, str]
 
 
+class PipInstallReport(TypedDict, total=False):
+    version: str
+    install: List[Dict[str, Any]]
+    environment: Dict[str, str]
+
+
+def _install_report_to_installed_distributions(
+    report: PipInstallReport,
+) -> InstalledDistributions:
+    environment = report["environment"]
+    dists = [
+        PipInstallReportItemInstalledDistribution(json_dist, environment)
+        for json_dist in report["install"]
+    ]
+    return {dist.name: dist for dist in dists}
+
+
+def pip_dry_run_install_project(
+    python: str,
+    constraints_filename: Path,
+    project_root: Path,
+    extras: Optional[Sequence[NormalizedName]] = None,
+) -> InstalledDistributions:
+    # dry-run install project with constraints
+    with tempfile.NamedTemporaryFile() as report:
+        report.close()
+        project_name = get_project_name(python, project_root)
+        log_info(f"Dry-run installing {project_name}")
+        cmd = [
+            python,
+            "-m",
+            "pip",
+            "install",
+            "--dry-run",
+            "--ignore-installed",
+            "--report",
+            report.name,
+            "-c",
+            f"{constraints_filename}",
+        ]
+        cmd.append("-e")
+        if extras:
+            extras_str = ",".join(extras)
+            cmd.append(f"{project_root}[{extras_str}]")
+        else:
+            cmd.append(f"{project_root}")
+        log_debug(f"Running {shlex.join(cmd)}")
+        constraints = constraints_filename.read_text(encoding="utf-8").strip()
+        if constraints:
+            log_debug(f"with {constraints_filename}:")
+            log_debug(constraints)
+        else:
+            log_debug(f"with empty {constraints_filename}.")
+        check_call(cmd)
+        with open(report.name, encoding="utf-8") as json_file:
+            return _install_report_to_installed_distributions(json.load(json_file))
+
+
 def pip_upgrade_project(
     python: str,
     constraints_filename: Path,