pip-reinstall-test.py

import atexit
import subprocess
import sys
import tempfile
import textwrap
from pathlib import Path
import shlex

from rich import print

HERE = Path(__file__).parent

PYTHON_VERSION = "python3.8"

PIP_VERSIONS: list[list[str | Path]] = [
    # ["pip<22.1"],
    ["-e", HERE / ".." / "pip"]
]

REFS_TO_CACHE = [
    "git+https://github.com/pypa/pip-test-package@7d654e66c8fa7149c165ddeffa5b56bc06619458",  # 0.1.1 on same tag
    "git+https://github.com/pypa/pip-test-package@5547fa909e83df8bd743d3978d6667497983a4b7",  # 0.1.1 on master branch
    "git+https://github.com/pypa/pip-test-package@f1c1020ebac81f9aeb5c766ff7a772f709e696ee",  # 0.1.2 on same tag
]

REFS_TO_WHEELHOUSE = [
    "git+https://github.com/pypa/pip-test-package@0.1.1",
    "git+https://github.com/pypa/pip-test-package@0.1.2",
]

# fmt: off
# case_name: (install_req, reinstall_req, options)
CASES = {
    "vcs-different-ref-different-version": (
        ["git+https://github.com/pypa/pip-test-package@0.1.2"],
        ["git+https://github.com/pypa/pip-test-package@0.1.1"],
        [
            ([], True),
        ],
    ),
    "vcs-same-ref-same-version": (
        ["git+https://github.com/pypa/pip-test-package@0.1.2"],
        ["git+https://github.com/pypa/pip-test-package@0.1.2"],
        [
            ([], False),
            (["--upgrade"], True),
            (["--force-reinstall"], True),
        ],
    ),
    "vcs-different-ref-same-version-different-commit": (
        ["pip-test-package @ git+https://github.com/pypa/pip-test-package@0.1.1"],
        ["pip-test-package @ git+https://github.com/pypa/pip-test-package@5547fa909e83df8bd743d3978d6667497983a4b7"],
        [
            ([], True),
        ],
    ),
    "vcs-different-ref-same-version-different-commit-no-cache": (
        ["--no-cache", "pip-test-package @ git+https://github.com/pypa/pip-test-package@0.1.1"],
        ["--no-cache","pip-test-package @ git+https://github.com/pypa/pip-test-package@5547fa909e83df8bd743d3978d6667497983a4b7"],
        [
            ([], True),
        ],
    ),
    "vcs-different-ref-same-commit": (
        ["pip-test-package @ git+https://github.com/pypa/pip-test-package@0.1.1"],
        ["pip-test-package @ git+https://github.com/pypa/pip-test-package@7d654e66c8fa7149c165ddeffa5b56bc06619458"],
        [
            ([], False),
            (["--upgrade"], True),  # TODO could avoid reinstall since same commit
            (["--force-reinstall"], True),
        ],
    ),
    "vcs-same-commit": (
        ["pip-test-package @ git+https://github.com/pypa/pip-test-package@7d654e66c8fa7149c165ddeffa5b56bc06619458"],
        ["pip-test-package @ git+https://github.com/pypa/pip-test-package@7d654e66c8fa7149c165ddeffa5b56bc06619458"],
        [
            ([], False),
            (["--upgrade"], False),  # TODO currently reinstall, could be avoided since same commit hash
            (["--force-reinstall"], True),
        ],
    ),
    "local-wheelhouse-no-version": (
        ["pip-test-package"],
        ["pip-test-package"],
        [
            ([], False),
            (["--upgrade"], False),
            (["--force-reinstall"], True),
        ],
    ),
    "local-wheelhouse-version-before-no-version-after": (
        ["pip-test-package==0.1.1"],
        ["pip-test-package"],
        [
            ([], False),
            (["--upgrade"], True)
        ],
    ),
    "local-wheelhouse-version-before-different-version-after": (
        ["pip-test-package==0.1.1"],
        ["pip-test-package==0.1.2"],
        [
            ([], True),
        ],
    ),
    "version-req-afer-vcs-ref": (
        ["pip-test-package @ git+https://github.com/pypa/pip-test-package@0.1.1"],
        ["pip-test-package==0.1.1"],
        [
            ([], True),  # TODO this currently does not reinstall
        ],
    ),
    "vcs-ref-afer-version-req": (
        ["pip-test-package==0.1.1"],
        ["pip-test-package @ git+https://github.com/pypa/pip-test-package@7d654e66c8fa7149c165ddeffa5b56bc06619458"],
        [
            ([], True),
        ],
    ),
    "file-wheel-same-version": (
        [HERE / "wheelhouse" / "pip_test_package-0.1.1-py3-none-any.whl" ],
        [HERE / "wheelhouse" / "pip_test_package-0.1.1-py3-none-any.whl" ],
        [
            ([], True),
        ],
    ),
    "file-sdist-same-version": (
        ["pip-test-package @ file://" + str(HERE / "wheelhouse" / "pip-test-package-0.1.1.tar.gz")],
        ["pip-test-package @ file://" + str(HERE / "wheelhouse" / "pip-test-package-0.1.1.tar.gz")],
        [
            ([], True),
        ],
    ),
}
# fmt: on


class PipInstallResult:
    def __init__(self, result: subprocess.CompletedProcess[str]):
        self.result = result

    @property
    def uninstalled(self) -> bool:
        return "Uninstalling pip-test-package" in self.result.stdout

    @property
    def installed(self) -> bool:
        return "Successfully installed pip-test-package" in self.result.stdout

    @property
    def reinstalled(self) -> bool:
        return self.installed and self.uninstalled

    @property
    def already_satisfied(self) -> bool:
        return (
            "Requirement already satisfied: pip-test-package" in self.result.stdout
            or (not self.installed and not self.reinstalled)
        )


class VirtualEnvironment:
    def __init__(self, venv_path: Path):
        self.venv_path = venv_path
        self.cache_path = HERE / "cache"
        self.wheelhouse_path = HERE / "wheelhouse"
        self.prepare_cache()
        self.prepare_wheelhouse()

    def prepare_cache(self) -> None:
        if self.cache_path.is_dir():
            return
        with tempfile.TemporaryDirectory() as tmpdir:
            for ref_to_cache in REFS_TO_CACHE:
                subprocess.check_call(
                    [
                        self.venv_path / "bin" / "pip",
                        "wheel",
                        "--no-index",
                        "--cache-dir",
                        self.cache_path,
                        "--wheel-dir",
                        tmpdir,
                        ref_to_cache,
                    ]
                )

    def prepare_wheelhouse(self) -> None:
        if self.wheelhouse_path.is_dir():
            return
        for ref_to_wheelhouse in REFS_TO_WHEELHOUSE:
            subprocess.check_call(
                [
                    self.venv_path / "bin" / "pip",
                    "wheel",
                    "--no-index",
                    "--no-cache",
                    "--wheel-dir",
                    self.wheelhouse_path,
                    ref_to_wheelhouse,
                ]
            )

    @classmethod
    def create(cls, pip_install_args: list[str | Path]) -> "VirtualEnvironment":
        tempdir = tempfile.TemporaryDirectory()
        venv_path = Path(tempdir.name)
        atexit.register(tempdir.cleanup)
        result = subprocess.run(
            ["virtualenv", venv_path, "-p", PYTHON_VERSION],
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
        assert result.returncode == 0, result.stdout
        venv = cls(venv_path)
        venv.pip_install("--upgrade", *pip_install_args)
        return venv

    def pip_install(
        self, *args: str | Path, env: dict[str, str] = {}
    ) -> subprocess.CompletedProcess[str]:
        result = subprocess.run(
            (self.venv_path / "bin" / "pip", "install") + args,  # type: ignore
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            env=env,
        )
        assert result.returncode == 0, result.stdout
        return result

    def pip_install_pip_test_package(self, *args: str | Path) -> PipInstallResult:
        result = self.pip_install(
            *args,
            env={
                "PIP_CACHE_DIR": str(self.cache_path),
                "PIP_NO_INDEX": "1",
                "PIP_FIND_LINKS": str(self.wheelhouse_path),
            },
        )
        return PipInstallResult(result)


def test_one(
    name: str,
    install_req: list[str],
    reinstall_req: list[str],
    reinstall_options: list[str],
    expect_reinstall: bool,
    print_output_always: bool,
) -> None:
    for pip_version in PIP_VERSIONS:
        venv = VirtualEnvironment.create(pip_version)
        print(f"[blue bold]{name} - pip {pip_version}")
        print("│ pip install", shlex.join(str(s) for s in install_req))
        r = venv.pip_install_pip_test_package(*install_req)
        assert not r.already_satisfied
        assert not r.uninstalled
        assert r.installed
        print("│ pip install", shlex.join(str(s) for s in (reinstall_options + reinstall_req)), end="")
        r = venv.pip_install_pip_test_package(*(reinstall_options + reinstall_req))
        print(f" > {r.already_satisfied=} {r.uninstalled=} {r.installed=}")
        success = not (r.reinstalled ^ expect_reinstall)
        if not success or print_output_always:
            print(
                "[dim]" + textwrap.indent(r.result.stdout, "[/dim]│[dim] ") + "[/dim]",
                end="",
            )
        if not success:
            if r.reinstalled:
                comment = "error: unexpected reinstall"
            else:
                comment = "error: did not reinstall"
        else:
            if expect_reinstall:
                comment = "ok: reinstalled, as expected"
            else:
                comment = "ok: did not reinstall, as expected"
        print("╰─>" + (f"[green] {comment}" if success else f"[red] {comment}"))


if len(sys.argv) > 1:
    name = sys.argv[1]
    install_req, reinstall_req, options = CASES[name]
    for reinstall_options, expect_reinstall in options:
        test_one(
            name,
            install_req,
            reinstall_req,
            reinstall_options,
            expect_reinstall,
            print_output_always=True,
        )
else:
    for name, (install_req, reinstall_req, options) in CASES.items():
        for reinstall_options, expect_reinstall in options:
            test_one(
                name,
                install_req,
                reinstall_req,
                reinstall_options,
                expect_reinstall,
                print_output_always=False,
            )