Venv

# -*- coding: utf-8 -*-

import os
from typing import Tuple
from configparser import ConfigParser

PYVENV_CFG_FILENAME = "pyvenv.cfg"
CFG_KEY_HOME = "home"
CFG_KEY_INCLUDE_SYSTEM = "include-system-site-packages"
CFG_KEY_VERSION = "version"

_FAKE_SECTION_NAME = "default"
_FAKE_DEFAULT_SECTION = f"[{_FAKE_SECTION_NAME}]"


class PyvenvCfg:

    home: str
    include_system_site_packages: bool
    version: str

    def to_version_tuple(self) -> Tuple[int, ...]:
        return tuple([int(d) for d in self.version.split("-")[0].split(".")])


def exists_pyvenv_cfg(env_root: str) -> bool:
    return os.path.exists(os.path.join(env_root, PYVENV_CFG_FILENAME))


def read_pyvenv_cfg(env_root: str) -> PyvenvCfg:
    if not env_root:
        raise ValueError("Empty argument")

    cfg_path = os.path.join(env_root, PYVENV_CFG_FILENAME)
    if not os.path.isfile(cfg_path):
        raise FileNotFoundError(f"Not found '{cfg_path}' file")
    if not os.access(cfg_path, os.R_OK):
        raise PermissionError(f"Not readable '{cfg_path}' file")

    with open(cfg_path) as f:
        content = _FAKE_DEFAULT_SECTION + "\n" + f.read()

    parser = ConfigParser()
    parser.read_string(content, cfg_path)
    home = parser.get(_FAKE_SECTION_NAME, CFG_KEY_HOME)
    include_system = parser.getboolean(_FAKE_SECTION_NAME, CFG_KEY_INCLUDE_SYSTEM)
    version = parser.get(_FAKE_SECTION_NAME, CFG_KEY_VERSION)

    result = PyvenvCfg()
    result.home = home
    result.include_system_site_packages = include_system
    result.version = version
    return result


def read_site_packages_dir(env_root: str) -> str:
    cfg = read_pyvenv_cfg(env_root)
    versions = cfg.to_version_tuple()
    if len(versions) < 2:
        raise RuntimeError(f"Wrong python version: {versions}")
    major = versions[0]
    minor = versions[1]
    return os.path.join(env_root, "lib", f"python{major}.{minor}", "site-packages")

# -*- coding: utf-8 -*-

import os
import tempfile
from unittest import TestCase, main
from recc.venv.pyvenv_cfg import read_pyvenv_cfg, read_site_packages_dir

_SAMPLE_PYENV_CFG = """home = /unknown/venv/root
include-system-site-packages = true
version = 3.7.3
"""


class PyvenvCfgTestCase(TestCase):
    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.pyenv_cfg_path = os.path.join(self.temp_dir.name, "pyvenv.cfg")
        with open(self.pyenv_cfg_path, "w") as f:
            f.write(_SAMPLE_PYENV_CFG)
        self.assertTrue(os.path.isfile(self.pyenv_cfg_path))

    def tearDown(self):
        self.temp_dir.cleanup()

    def test_read_pyvenv_cfg(self):
        cfg = read_pyvenv_cfg(self.temp_dir.name)
        self.assertEqual("/unknown/venv/root", cfg.home)
        self.assertTrue(cfg.include_system_site_packages)
        self.assertEqual("3.7.3", cfg.version)

    def test_read_site_packages_dir(self):
        path = read_site_packages_dir(self.temp_dir.name)
        self.assertTrue(path.endswith("lib/python3.7/site-packages"))


if __name__ == "__main__":
    main()

# -*- coding: utf-8 -*-

import sys
import os
from types import SimpleNamespace
from typing import Optional
from venv import EnvBuilder
from shutil import rmtree
from recc.subprocess.async_python_subprocess import AsyncPythonSubprocess
from recc.venv.venv_context_changer import VenvContextChanger


def _get_site_packages_dir(env_dir) -> str:
    if sys.platform == "win32":
        return os.path.join(env_dir, "Lib", "site-packages")
    else:
        return os.path.join(
            env_dir, "lib", "python%d.%d" % sys.version_info[:2], "site-packages"
        )


class AsyncVirtualEnvironment:
    def __init__(
        self,
        root_directory: str,
        system_site_packages=False,
        pip_timeout: Optional[float] = None,
        *,
        isolate_ensure_pip=True,
    ):
        self._root_directory = root_directory
        self._pip_timeout = pip_timeout if pip_timeout else 0.0
        self._isolate_ensure_pip = isolate_ensure_pip
        self._venv = EnvBuilder(
            system_site_packages=system_site_packages,
            clear=False,
            symlinks=False,
            upgrade=False,
            with_pip=True,
            prompt=None,
        )

        self._context = self._venv.ensure_directories(os.path.abspath(root_directory))
        if not hasattr(self._context, "pip_exe"):
            self._context.pip_exe = os.path.join(
                self._context.bin_path, "pip%d.%d" % sys.version_info[:2]
            )
        if not hasattr(self._context, "site_packages_dir"):
            self._context.site_packages_dir = _get_site_packages_dir(root_directory)

    @property
    def root(self) -> str:
        """
        The original venv root directory.

        Preserves the original value passed to the constructor argument.
        """
        return self._root_directory

    @property
    def context(self) -> SimpleNamespace:
        return self._context

    @property
    def bin_name(self) -> str:
        return self.context.bin_name

    @property
    def bin_path(self) -> str:
        return self.context.bin_path

    @property
    def env_dir(self) -> str:
        """
        Absolute path to the venv root directory.
        """
        return self.context.env_dir

    @property
    def env_exe(self) -> str:
        return self.context.env_exe

    @property
    def env_name(self) -> str:
        return self.context.env_name

    @property
    def executable(self) -> str:
        return self.context.executable

    @property
    def inc_path(self) -> str:
        return self.context.inc_path

    @property
    def prompt(self) -> str:
        return self.context.prompt

    @property
    def python_dir(self) -> str:
        return self.context.python_dir

    @property
    def python_exe(self) -> str:
        return self.context.python_exe

    @property
    def pip_exe(self) -> str:
        return self._context.pip_exe

    @property
    def site_packages_dir(self) -> str:
        return self._context.site_packages_dir

    @property
    def exists(self) -> bool:
        return os.path.exists(self.env_exe)

    def _create(self) -> None:
        self._venv.create(self._root_directory)

    def _create_if_not_exists(self) -> None:
        if not self.exists:
            self._venv.create(self._root_directory)

    def _clear(self) -> None:
        self._venv.clear_directory(self._root_directory)

    def _create_configuration(self) -> None:
        self._venv.create_configuration(self._context)

    def _setup_python(self) -> None:
        self._venv.setup_python(self._context)

    def _setup_scripts(self) -> None:
        self._venv.setup_scripts(self._context)

    def _upgrade_dependencies(self) -> None:
        func = getattr(self._venv, "upgrade_dependencies")
        if func is None:
            raise NotImplementedError
        assert sys.version_info >= (3, 9)
        func(self._context)

    def _post_setup(self) -> None:
        self._venv.post_setup(self._context)

    def _install_scripts(self, path: str) -> None:
        self._venv.install_scripts(self._context, path)

    async def _setup_pip(self) -> None:
        await AsyncPythonSubprocess(self.env_exe).ensure_pip(
            isolate=self._isolate_ensure_pip
        )

    async def create(self) -> None:
        # See issue 24875. We need system_site_packages to be False
        # until after pip is installed.
        true_system_site_packages = self._venv.system_site_packages
        self._venv.system_site_packages = False
        self._create_configuration()
        self._setup_python()
        if self._venv.with_pip:
            await self._setup_pip()
        if not self._venv.upgrade:
            self._setup_scripts()
            self._post_setup()
        if true_system_site_packages:
            # We had set it to False before, now
            # restore it and rewrite the configuration
            self._venv.system_site_packages = True
            self._create_configuration()

    async def create_if_not_exists(self, remove_if_raised=True) -> None:
        if not self.exists:
            try:
                await self.create()
            except BaseException as e:
                if remove_if_raised:
                    rmtree(self.env_dir, ignore_errors=True)
                raise RuntimeError(
                    f"Failed to create virtual environment: '{self.env_dir}'"
                ) from e

    def create_python_subprocess(self) -> AsyncPythonSubprocess:
        return AsyncPythonSubprocess(self.env_exe, self._pip_timeout)

    def create_context_changer(self):
        return VenvContextChanger(
            env_dir=self.env_dir,
            env_exe=self.env_exe,
            bin_path=self.bin_path,
            site_packages_dir=self.site_packages_dir,
        )

# -*- coding: utf-8 -*-

import os
import sys
import site
from abc import abstractmethod
from typing import List, Optional
from overrides import overrides

real_prefix = "real_prefix"

X = sys.version_info[0]
Y = sys.version_info[1]
XY = f"{X}{Y}"
X_Y = f"{X}.{Y}"

LIB_PYTHON_ZIP = f"{sys.base_prefix}/{sys.platlibdir}/python{XY}.zip"
LIB_PYTHON = f"{sys.base_prefix}/{sys.platlibdir}/python{X_Y}"
LIB_PYTHON_LIB_DYNLOAD = f"{sys.base_prefix}/{sys.platlibdir}/python{X_Y}/lib-dynload"

PATH = "PATH"
VIRTUAL_ENV = "VIRTUAL_ENV"


class ContextChanger:
    @abstractmethod
    def open(self) -> None:
        raise NotImplementedError

    @abstractmethod
    def close(self) -> None:
        raise NotImplementedError

    def __enter__(self):
        self.open()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    async def __aenter__(self):
        raise RuntimeError(
            "Accessing the sys package while the event loop is running causes problems"
        )

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        raise RuntimeError(
            "Accessing the sys package while the event loop is running causes problems"
        )


class FakeContextChanger(ContextChanger):
    @overrides
    def open(self) -> None:
        pass

    @overrides
    def close(self) -> None:
        pass


class VenvContextChanger(ContextChanger):
    """
    References:
        - https://docs.python.org/3/library/sys.html
        - https://docs.python.org/3/library/venv.html
        - https://peps.python.org/pep-0405/
    """

    _sys_prefix: str
    _sys_exec_prefix: str
    _sys_executable: str
    _sys_path: List[str]
    _sys_real_prefix: Optional[str]
    _env_path: Optional[str]
    _env_virtual_env: Optional[str]

    def __init__(
        self,
        env_dir: str,
        env_exe: str,
        bin_path: str,
        site_packages_dir: str,
    ):
        self.env_dir = env_dir
        self.env_exe = env_exe
        self.bin_path = bin_path
        self.site_packages_dir = site_packages_dir

    @overrides
    def open(self) -> None:
        self._sys_prefix = sys.prefix
        self._sys_exec_prefix = sys.exec_prefix
        self._sys_executable = sys.executable
        self._sys_path = sys.path

        sys.prefix = self.env_dir
        sys.exec_prefix = self.env_dir
        sys.executable = self.env_exe
        sys.path = [
            LIB_PYTHON_ZIP,
            LIB_PYTHON,
            LIB_PYTHON_LIB_DYNLOAD,
        ]
        site.addsitedir(self.site_packages_dir)

        self._sys_real_prefix = getattr(sys, real_prefix, None)
        setattr(sys, real_prefix, self.env_dir)

        self._env_path = os.environ.get(PATH, None)
        self._env_virtual_env = os.environ.get(VIRTUAL_ENV, None)

        os.environ[PATH] = os.pathsep.join(
            [self.bin_path] + os.environ.get(PATH, "").split(os.pathsep)
        )
        os.environ[VIRTUAL_ENV] = self.env_dir

    @staticmethod
    def _restore_env(key: str, original: Optional[str]):
        if original is not None:
            os.environ[key] = original
        else:
            os.environ.pop(key)

    @overrides
    def close(self) -> None:
        sys.prefix = self._sys_prefix
        sys.exec_prefix = self._sys_exec_prefix
        sys.executable = self._sys_executable
        sys.path = self._sys_path

        assert hasattr(sys, real_prefix)
        if self._sys_real_prefix is not None:
            setattr(sys, real_prefix, self._sys_real_prefix)
        else:
            delattr(sys, real_prefix)

        self._restore_env(PATH, self._env_path)
        self._restore_env(VIRTUAL_ENV, self._env_virtual_env)