Rewrite virtualenv tool in tests for 20+ support

Co-Authored-By: Lumir Balhar <[email protected]>

Author: uranusjr

tests/lib/venv.py

@@ -1,8 +1,7 @@
 import compileall
 import os
 import shutil
-import subprocess
-import sys
+import sysconfig
 import textwrap
 import venv as _venv
 from pathlib import Path
@@ -47,15 +46,16 @@ def __init__(
         self._create()
 
     def _update_paths(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)
+        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}>"
@@ -64,29 +64,21 @@ 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 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._venv_type == "virtualenv":
-                subprocess.check_call(
+                _virtualenv.cli_run(
                     [
-                        sys.executable,
-                        "-m",
-                        "virtualenv",
                         "--no-pip",
                         "--no-wheel",
                         "--no-setuptools",
-                        str(self.location),
-                    ]
+                        os.fspath(self.location),
+                    ],
                 )
-                self._fix_virtualenv_site_module()
             elif self._venv_type == "venv":
                 builder = _venv.EnvBuilder()
                 context = builder.ensure_directories(self.location)
@@ -96,71 +88,30 @@ def _create(self, clear: bool = False) -> None:
             self.sitecustomize = self._sitecustomize
             self.user_site_packages = self._user_site_packages
 
-    def _fix_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:
-        contents = ""
-        if self._venv_type == "venv":
-            # Enable user site (before system).
-            contents += textwrap.dedent(
-                """
-                import os, site, sys
-
-                if not os.environ.get('PYTHONNOUSERSITE', False):
-
-                    site.ENABLE_USER_SITE = True
-
-                    # 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.
-                    site.addsitedir(site.getusersitepackages())
-
-                    # Third, add back system-sites related paths.
-                    for path in site.getsitepackages():
-                        site.addsitedir(path)
-                """
-            ).strip()
+        # Enable user site (before system).
+        contents = textwrap.dedent(
+            """
+            import os, site, sys
+            if not os.environ.get('PYTHONNOUSERSITE', False):
+                site.ENABLE_USER_SITE = True
+                # 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.
+                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"
@@ -191,12 +142,12 @@ def user_site_packages(self) -> bool:
 
     @user_site_packages.setter
     def user_site_packages(self, value: bool) -> None:
-        self._user_site_packages = value
-        if self._venv_type == "virtualenv":
-            marker = self.lib / "no-global-site-packages.txt"
-            if self._user_site_packages:
-                marker.unlink()
-            else:
-                marker.touch()
-        elif self._venv_type == "venv":
-            self._customize_site()
+        self._customize_site()
+        pyvenv_cfg = self.location.joinpath("pyvenv.cfg")
+        modified_lines = []
+        for line in pyvenv_cfg.read_text().splitlines():
+            k, v = line.split("=", 1)
+            if k.strip() == "include-system-site-packages":
+                line = f"include-system-site-packages = {str(bool(value)).lower()}"
+            modified_lines.append(line)
+        pyvenv_cfg.write_text("\n".join(modified_lines))