pip/tests/lib/venv.py

255 lines
9.3 KiB
Python

import compileall
import os
import shutil
import subprocess
import sys
import sysconfig
import textwrap
import venv as _venv
from pathlib import Path
from typing import TYPE_CHECKING, Dict, Optional, Union
import virtualenv as _virtualenv
if TYPE_CHECKING:
# Literal was introduced in Python 3.8.
from typing import Literal
VirtualEnvironmentType = Literal["virtualenv", "venv"]
else:
VirtualEnvironmentType = str
class VirtualEnvironment:
"""
An abstraction around virtual environments, currently it only uses
virtualenv but in the future it could use pyvenv.
"""
def __init__(
self,
location: Path,
template: Optional["VirtualEnvironment"] = None,
venv_type: Optional[VirtualEnvironmentType] = None,
) -> None:
self.location = location
assert template is None or venv_type is None
self._venv_type: VirtualEnvironmentType
if template is not None:
self._venv_type = template._venv_type
elif venv_type is not None:
self._venv_type = venv_type
else:
self._venv_type = "virtualenv"
self._user_site_packages = False
self._template = template
self._sitecustomize: Optional[str] = None
self._update_paths()
self._create()
@property
def _legacy_virtualenv(self) -> bool:
if self._venv_type != "virtualenv":
return False
return int(_virtualenv.__version__.split(".", 1)[0]) < 20
def __update_paths_legacy(self) -> None:
home, lib, inc, bin = _virtualenv.path_locations(self.location)
self.bin = Path(bin)
self.site = Path(lib) / "site-packages"
# Workaround for https://github.com/pypa/virtualenv/issues/306
if hasattr(sys, "pypy_version_info"):
version_dir = str(sys.version_info.major)
self.lib = Path(home, "lib-python", version_dir)
else:
self.lib = Path(lib)
def _update_paths(self) -> None:
if self._legacy_virtualenv:
self.__update_paths_legacy()
return
bases = {
"installed_base": self.location,
"installed_platbase": self.location,
"base": self.location,
"platbase": self.location,
}
paths = sysconfig.get_paths(vars=bases)
self.bin = Path(paths["scripts"])
self.site = Path(paths["purelib"])
self.lib = Path(paths["stdlib"])
def __repr__(self) -> str:
return f"<VirtualEnvironment {self.location}>"
def _create(self, clear: bool = False) -> None:
if clear:
shutil.rmtree(self.location)
if self._template:
# On Windows, calling `_virtualenv.path_locations(target)`
# will have created the `target` directory...
if (
self._legacy_virtualenv
and sys.platform == "win32"
and self.location.exists()
):
self.location.rmdir()
# Clone virtual environment from template.
shutil.copytree(self._template.location, self.location, symlinks=True)
self._sitecustomize = self._template.sitecustomize
self._user_site_packages = self._template.user_site_packages
else:
# Create a new virtual environment.
if self._legacy_virtualenv:
subprocess.check_call(
[
sys.executable,
"-m",
"virtualenv",
"--no-pip",
"--no-wheel",
"--no-setuptools",
os.fspath(self.location),
]
)
self._fix_legacy_virtualenv_site_module()
elif self._venv_type == "virtualenv":
_virtualenv.cli_run(
[
"--no-pip",
"--no-wheel",
"--no-setuptools",
os.fspath(self.location),
],
)
elif self._venv_type == "venv":
builder = _venv.EnvBuilder()
context = builder.ensure_directories(os.fspath(self.location))
builder.create_configuration(context)
builder.setup_python(context)
self.site.mkdir(parents=True, exist_ok=True)
else:
raise RuntimeError(f"Unsupported venv type {self._venv_type!r}")
self.sitecustomize = self._sitecustomize
self.user_site_packages = self._user_site_packages
def _fix_legacy_virtualenv_site_module(self) -> None:
# Patch `site.py` so user site work as expected.
site_py = self.lib / "site.py"
with open(site_py) as fp:
site_contents = fp.read()
for pattern, replace in (
(
# Ensure enabling user site does not result in adding
# the real site-packages' directory to `sys.path`.
("\ndef virtual_addsitepackages(known_paths):\n"),
(
"\ndef virtual_addsitepackages(known_paths):\n"
" return known_paths\n"
),
),
(
# Fix sites ordering: user site must be added before system.
(
"\n paths_in_sys = addsitepackages(paths_in_sys)"
"\n paths_in_sys = addusersitepackages(paths_in_sys)\n"
),
(
"\n paths_in_sys = addusersitepackages(paths_in_sys)"
"\n paths_in_sys = addsitepackages(paths_in_sys)\n"
),
),
):
assert pattern in site_contents
site_contents = site_contents.replace(pattern, replace)
with open(site_py, "w") as fp:
fp.write(site_contents)
# Make sure bytecode is up-to-date too.
assert compileall.compile_file(str(site_py), quiet=1, force=True)
def _customize_site(self) -> None:
if self._legacy_virtualenv:
contents = ""
else:
# Enable user site (before system).
contents = textwrap.dedent(
f"""
import os, site, sys
if not os.environ.get('PYTHONNOUSERSITE', False):
site.ENABLE_USER_SITE = {self._user_site_packages}
# First, drop system-sites related paths.
original_sys_path = sys.path[:]
known_paths = set()
for path in site.getsitepackages():
site.addsitedir(path, known_paths=known_paths)
system_paths = sys.path[len(original_sys_path):]
for path in system_paths:
if path in original_sys_path:
original_sys_path.remove(path)
sys.path = original_sys_path
# Second, add user-site.
if {self._user_site_packages}:
site.addsitedir(site.getusersitepackages())
# Third, add back system-sites related paths.
for path in site.getsitepackages():
site.addsitedir(path)
"""
).strip()
if self._sitecustomize is not None:
contents += "\n" + self._sitecustomize
sitecustomize = self.site / "sitecustomize.py"
sitecustomize.write_text(contents)
# Make sure bytecode is up-to-date too.
assert compileall.compile_file(str(sitecustomize), quiet=1, force=True)
def _rewrite_pyvenv_cfg(self, replacements: Dict[str, str]) -> None:
pyvenv_cfg = self.location.joinpath("pyvenv.cfg")
lines = pyvenv_cfg.read_text(encoding="utf-8").splitlines()
def maybe_replace_line(line: str) -> str:
key = line.split("=", 1)[0].strip()
try:
value = replacements[key]
except KeyError: # No need to replace.
return line
return f"{key} = {value}"
lines = [maybe_replace_line(line) for line in lines]
pyvenv_cfg.write_text("\n".join(lines), encoding="utf-8")
def clear(self) -> None:
self._create(clear=True)
def move(self, location: Union[Path, str]) -> None:
shutil.move(os.fspath(self.location), location)
self.location = Path(location)
self._update_paths()
@property
def sitecustomize(self) -> Optional[str]:
return self._sitecustomize
@sitecustomize.setter
def sitecustomize(self, value: str) -> None:
self._sitecustomize = value
self._customize_site()
@property
def user_site_packages(self) -> bool:
return self._user_site_packages
@user_site_packages.setter
def user_site_packages(self, value: bool) -> None:
self._user_site_packages = value
if self._legacy_virtualenv:
marker = self.lib / "no-global-site-packages.txt"
if self._user_site_packages:
marker.unlink()
else:
marker.touch()
else:
self._rewrite_pyvenv_cfg(
{"include-system-site-packages": str(bool(value)).lower()}
)
self._customize_site()