venv.py

import ast
import json
import os
import shlex
import shutil
import sys
import textwrap
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any

import virtualenv

from pyodide_build.build_env import (
    get_build_flag,
    get_pyodide_root,
    in_xbuildenv,
    wheel_platform,
)
from pyodide_build.common import IS_WIN, run_command
from pyodide_build.logger import logger

# A subset of supported virtualenv options that make sense in Pyodide's context.
# Our aim will not be to support all of them, and some of them will never be in
# the list, for example, --no-pip and so on. We provide these on a best-effort
# basis as they should work and are easy to test.
SUPPORTED_VIRTUALENV_OPTIONS = [
    "--clear",
    "--no-clear",
    "--no-vcs-ignore",
    # "--copies", "--always-copy", FIXME: node fails to invoke Pyodide
    # "--symlink-app-data", FIXME: node fails to invoke Pyodide
    "--no-download",
    "--never-download",
    "--download",
    "--extra-search-dir",
    "--pip",
    "--setuptools",
    "--no-setuptools",
    "--no-periodic-update",
]

# See https://github.com/python/cpython/issues/119535
# fmt: off
PI_ALIASES = (
    "pythonπ",           # pythonπ U+03C0 GREEK SMALL LETTER PI
    "python\U0001d70b",  # python𝜋 U+1D70B MATHEMATICAL ITALIC SMALL PI
    "πthon",             # πthon
    "\U0001d70bthon",    # 𝜋thon
)
# fmt: on


def dedent(s: str) -> str:
    return textwrap.dedent(s).strip() + "\n"


def _pip_script_name(pip: Path, exe_suffix: str) -> Path:
    """Return the pip script path with *exe_suffix* applied.

    On Unix ``exe_suffix`` is empty and the name is returned unchanged, so
    versioned names like ``pip3.12`` keep their version (``Path.with_suffix("")``
    would wrongly strip it to ``pip3``). On Windows the real extension (``.exe``
    or ``.bat``) is replaced with ``exe_suffix``.
    """
    if not exe_suffix:
        return pip
    base_name = pip.stem if pip.suffix else pip.name
    return pip.parent / (base_name + exe_suffix)


def get_pyversion() -> str:
    return f"{sys.version_info.major}.{sys.version_info.minor}"


def check_host_python_version(session: Any) -> tuple[int, int]:
    pyodide_version_str = session.interpreter.version.partition(" ")[0].split(".")[:2]
    pyodide_version = (int(pyodide_version_str[0]), int(pyodide_version_str[1]))
    sys_version = (sys.version_info.major, sys.version_info.minor)
    if pyodide_version == sys_version:
        return pyodide_version
    pyodide_version_fmt = ".".join(pyodide_version_str)
    sys_version_fmt = ".".join(str(version) for version in sys_version)
    logger.stderr(
        f"Expected host Python version to be {pyodide_version_fmt} but got version {sys_version_fmt}"
    )
    sys.exit(1)


def pyodide_dist_dir() -> Path:
    return get_pyodide_root() / "dist"


class PyodideVenv(ABC):
    """Base class for creating Pyodide virtual environments.

    This class contains common functionality shared across all platforms.
    Platform-specific implementations should inherit from this class.
    """

    def __init__(self, dest: Path, virtualenv_args: list[str] | None = None) -> None:
        self.dest = dest
        self.virtualenv_args = virtualenv_args or []
        self._venv_root: Path | None = None
        self._venv_bin: Path | None = None
        self._pyodide_pyversion: tuple[int, int] | None = None

    @property
    def venv_root(self) -> Path:
        """Get the path to the virtualenv's root directory."""
        if self._venv_root is None:
            raise RuntimeError("venv_root is not set")

        return self._venv_root

    @property
    def venv_bin(self) -> Path:
        """Get the path to the virtualenv's bin directory."""
        if self._venv_bin is None:
            raise RuntimeError("venv_bin is not set")

        return self._venv_bin

    @property
    def exe_suffix(self) -> str:
        """Return the executable suffix for the platform."""
        return ""

    @property
    def python_exe_name(self) -> str:
        """Return the Python executable name for the platform."""
        return "python" + self.exe_suffix

    @property
    @abstractmethod
    def bin_dir_name(self) -> str:
        """Return the bin directory name for the platform."""
        pass

    @property
    def pyodide_exe_name(self) -> str:
        """Return the Pyodide executable name for the platform."""
        return "pyodide" + self.exe_suffix

    @property
    def host_python_name(self) -> str:
        """Get the host python executable name."""
        return f"python{get_pyversion()}-host" + self.exe_suffix

    @property
    def host_python_name_noversion(self) -> str:
        """Get the host python executable name without version."""
        return "python-host" + self.exe_suffix

    @property
    @abstractmethod
    def host_python_symlink_suffix(self) -> str:
        """Get the host python symlink suffix."""
        pass

    @property
    @abstractmethod
    def host_sitepackages_path(self) -> Path:
        """
        Path to the host-specific site-packages directory in the virtualenv
        """
        pass

    @property
    def venv_sitepackages_path(self) -> Path:
        """
        Path to the site-packages directory in the virtualenv, where packages should be installed.

        Note that in host environment, Windows uses 'Lib\\site-packages' while Unix uses 'lib/pythonX.Y/site-packages'.
        However, Pyodide environment is Unix-like, so we always use the Unix-style path here so that packages are located correctly
        inside the Pyodide virtual environment
        """
        return self.venv_root / "lib" / f"python{get_pyversion()}" / "site-packages"

    @property
    def interpreter_path(self) -> Path:
        """Get the path to the original Pyodide Python interpreter."""
        return pyodide_dist_dir() / self.python_exe_name

    @property
    def interpreter_symlink_path(self) -> Path:
        """Get the path to the Pyodide Python interpreter symlink."""
        return self.venv_bin / self.python_exe_name

    @property
    def script_interpreter_suffix(self) -> str:
        """Get the extension of the interpreter that console scripts invoke.

        This is the ``+ suffix`` half of the calculation that ``pip_wrapper``
        does to turn ``sys.executable`` into the Pyodide interpreter, the other
        half being ``host_python_symlink_suffix``.
        """
        return ""

    @property
    def script_interpreter_path(self) -> Path:
        """Get the path to the interpreter that installed console scripts invoke."""
        return self.venv_bin / f"python{self.script_interpreter_suffix}"

    @property
    def pyodide_cli_path(self) -> Path:
        """Get the path to the pyodide CLI script in the virtualenv."""
        return self.venv_bin / self.pyodide_exe_name

    @property
    def host_python_path(self) -> Path:
        """Get the path to the host python executable in the virtualenv."""
        return self.venv_bin / self.host_python_name

    @property
    def host_python_path_noversion(self) -> Path:
        """Get the path to the host python executable without version in the virtualenv."""
        return self.venv_bin / self.host_python_name_noversion

    @property
    def host_python_symlink_path(self) -> Path:
        """Get the path to the host python symlink in the virtualenv."""
        return self.venv_bin / f"python{self.host_python_symlink_suffix}"

    @property
    def pip_conf_path(self) -> Path:
        """Get the path to the pip.conf file in the virtualenv."""
        return self.venv_root / "pip.conf"

    @property
    def pip_patched_path(self) -> Path:
        """Get the path to the pip_patched script in the virtualenv."""
        return self.venv_bin / "pip_patched"

    @property
    def pip_wrapper_path(self) -> Path:
        """Get the path to the pip wrapper script in the virtualenv."""
        return self.venv_bin / "_pip-wrapper.py"

    @property
    @abstractmethod
    def host_python_wrapper(self) -> str:
        """Get the content of the host python wrapper script.
        This script allows invoking the host python with the correct PYTHONHOME.
        """
        pass

    @property
    @abstractmethod
    def host_pip_wrapper(self) -> str:
        """Get the content of the host pip wrapper script."""
        pass

    def validate_interpreter(self) -> None:
        """Validate that the Pyodide interpreter exists."""
        if not self.interpreter_path.exists():
            raise RuntimeError(
                f"Pyodide python interpreter not found at {self.interpreter_path}"
            )

    def get_cli_args(self) -> list[str]:
        """Build the CLI arguments for virtualenv."""
        cli_args = ["--python", str(self.interpreter_path)]

        if self.virtualenv_args:
            for arg in self.virtualenv_args:
                if arg.startswith("--"):
                    arg_name = arg.split("=")[0] if "=" in arg else arg
                    if arg_name not in SUPPORTED_VIRTUALENV_OPTIONS:
                        msg = f"Unsupported virtualenv option: {arg_name}"
                        logger.warning(msg)

            cli_args.extend(self.virtualenv_args)

        return cli_args

    def _create_pip_conf(self) -> None:
        """Create pip.conf file in venv root.

        This file adds a few options that will always be used by pip install.
        """
        if in_xbuildenv():
            # In the xbuildenv, we don't have the packages locally. We will include
            # in the xbuildenv a PEP 503 index for the vendored Pyodide packages
            # https://peps.python.org/pep-0503/
            repo = f"extra-index-url={(get_pyodide_root() / 'package_index').as_uri()}"
        else:
            # In the Pyodide development environment, the Pyodide dist directory
            # should contain the needed wheels. find-links
            repo = f"find-links={pyodide_dist_dir()}"

        # Prevent attempts to install binary wheels from source.
        # Maybe some day we can convince pip to invoke `pyodide build` as the build
        # front end for wheels...
        self.pip_conf_path.write_text(
            dedent(
                f"""
                [global]
                only-binary=:all:
                {repo}
                """
            )
        )

    def _install_stdlib(self) -> None:
        """Install micropip and all unvendored stdlib modules."""
        logger.info("... Installing standard library")

        # Micropip we could install with pip hypothetically, but because we use
        # `--extra-index-url` it would install the pypi version which we don't want.

        # Other stuff we need to load with loadPackage
        to_load = ["micropip"]
        run_command(
            [
                self.interpreter_symlink_path,
                "-c",
                dedent(
                    f"""
                    from pyodide_js import loadPackage
                    from pyodide_js._api import lockfile_packages
                    from pyodide_js._api import lockfile_unvendored_stdlibs_and_test
                    shared_libs = [pkgname for (pkgname,pkg) in lockfile_packages.object_entries() if getattr(pkg, "package_type", None) == "shared_library"]

                    to_load = [*lockfile_unvendored_stdlibs_and_test, *shared_libs, *{to_load!r}]
                    loadPackage(to_load);
                    """
                ),
            ],
            err_msg="ERROR: failed to install unvendored stdlib modules",
        )

    def _get_pyodide_pip_config(self) -> dict[str, Any]:
        """Compute the config used by the pip wrapper to emulate Pyodide's environment.

        The returned dict is serialized to ``pyodide_pip_config.json`` next to the
        pip wrapper script and loaded by ``pip_wrapper.py`` at runtime.
        """
        result = run_command(
            [
                self.interpreter_symlink_path,
                "-c",
                dedent(
                    """
                    import os, sys, sysconfig, platform;
                    print([
                        os.name,
                        sys.platform,
                        platform.system(),
                        sys.implementation._multiarch,
                        sysconfig.get_platform()
                    ])
                    """.replace(
                        "\n", ""
                    )  # Windows doesn't seems to like newlines here...
                ),
            ],
            err_msg="ERROR: failed to invoke Pyodide",
        )
        sysconfigdata_dir = str(
            Path(get_build_flag("TARGETINSTALLDIR")) / "sysconfigdata"
        )
        return dict(
            executable_symlink_suffix=self.host_python_symlink_suffix,
            exe_suffix=self.exe_suffix,
            script_interpreter_suffix=self.script_interpreter_suffix,
            pip_patched_name=self.pip_patched_path.name,
            pip_wrapper_name=self.pip_wrapper_path.name,
            platform_data=ast.literal_eval(result.stdout),
            pyodide_platform=wheel_platform(),
            sysconfigdata_dir=sysconfigdata_dir,
        )

    def _create_pip_script(self) -> None:
        """Create pip and write it into the virtualenv bin folder."""
        # pip needs to run in the host Python not in Pyodide, so we'll use the host
        # Python in the shebang. Use whichever Python was used to invoke
        # pyodide venv.

        # To support the "--clear" and "--no-clear" args, we need to remove
        # the existing symlinks before creating new ones.
        self.host_python_path.unlink(missing_ok=True)
        self.host_python_path_noversion.unlink(missing_ok=True)
        self.host_python_symlink_path.unlink(missing_ok=True)

        # Replace all pip* scripts in the venv bin folder with symlinks to
        # our patched pip script.
        for pip in self.venv_bin.glob("pip*"):
            if pip == self.pip_patched_path:
                continue
            pip.unlink(missing_ok=True)

            patched_pip_exe = _pip_script_name(pip, self.exe_suffix)
            if patched_pip_exe != self.pip_patched_path:
                patched_pip_exe.unlink(missing_ok=True)
                patched_pip_exe.symlink_to(self.pip_patched_path)

        # Weird hack to work around:
        # https://github.com/astral-sh/python-build-standalone/issues/380
        # If we resolve the symlink all the way, the python-host interpreter works
        # but won't install into our pyodide venv. If we don't resolve the symlink,
        # sys.prefix is calculated incorrectly. To ensure that we get the right
        # sys.prefix, we explicitly set it with the PYTHONHOME environment variable
        # and then call the symlink.
        self.host_python_symlink_path.symlink_to(sys.executable)
        self.host_python_path.write_text(self.host_python_wrapper)
        self.host_python_path.chmod(0o777)
        self.host_python_path_noversion.symlink_to(self.host_python_path)

        self.pip_patched_path.write_text(self.host_pip_wrapper)
        self.pip_patched_path.chmod(0o777)

        pip_wrapper_src = Path(__file__).parent / "pip_wrapper.py"
        shutil.copy(pip_wrapper_src, self.pip_wrapper_path)
        (self.venv_bin / "pyodide_pip_config.json").write_text(
            json.dumps(self._get_pyodide_pip_config())
        )

        # On windows, link the venv site-packages to the host site-packages so that the packages
        # installed in the windows location are visible to pyodide environment.
        # Note that if we simply symlink the host site-packages to the venv site-packages,
        # python doesn't seem to recognize the symlink as a real path, ending up not adding it to sys.path.
        # To work around this, we first move the host site-packages to the venv location,
        # then create a symlink from the host location to the venv location.
        if not self.venv_sitepackages_path.exists():
            self.venv_sitepackages_path.parent.mkdir(parents=True, exist_ok=True)
            self.host_sitepackages_path.replace(self.venv_sitepackages_path)
            self.host_sitepackages_path.symlink_to(
                self.venv_sitepackages_path, target_is_directory=True
            )

    @abstractmethod
    def _create_pyodide_script(self) -> None:
        """Create pyodide CLI script in the virtualenv bin folder.

        This is platform-specific and must be implemented by subclasses.
        """
        pass

    @abstractmethod
    def _create_python_symlink(self) -> None:
        """Create a symlink to the Pyodide Python interpreter."""
        pass

    @abstractmethod
    def _create_session(self):
        """Create and return a virtualenv session object."""
        pass

    def _create_pi_aliases(self) -> None:
        """
        Create pythonπ-style easter egg aliases for the Python 3.14
        release. See https://github.com/python/cpython/issues/119535 and
        https://github.com/python/cpython/pull/119536 for details. This
        is a no-op for all other versions of Python.
        """
        if self._pyodide_pyversion != (3, 14):
            return

        for name in PI_ALIASES:
            alias_path = self.venv_bin / f"{name}{self.exe_suffix}"
            alias_path.unlink(missing_ok=True)
            alias_path.symlink_to(self.interpreter_path)

    def configure_virtualenv(self) -> None:
        """Configure the virtualenv after creation."""
        logger.info("... Configuring virtualenv")
        self._create_python_symlink()
        self._create_pip_conf()
        self._create_pip_script()
        self._create_pyodide_script()
        self._create_pi_aliases()

    def create(self) -> None:
        """Create the Pyodide virtualenv."""
        logger.info("Creating Pyodide virtualenv at %s", self.dest)
        self.validate_interpreter()
        session = self._create_session()
        self._pyodide_pyversion = check_host_python_version(session)

        dest_path = Path(session.creator.dest)
        dest_existed_before = dest_path.exists()

        try:
            session.run()
            self._venv_root = dest_path.absolute()
            self._venv_bin = self._venv_root / self.bin_dir_name

            self.configure_virtualenv()
            self._install_stdlib()
        except (Exception, KeyboardInterrupt, SystemExit):
            if dest_existed_before:
                logger.warning(
                    "Venv creation failed; leaving pre-existing directory %s intact",
                    dest_path,
                )
            else:
                shutil.rmtree(dest_path, ignore_errors=True)
            raise

        logger.success("Successfully created Pyodide virtual environment!")


class UnixPyodideVenv(PyodideVenv):
    """Unix-specific implementation of Pyodide virtual environment creation."""

    @property
    def bin_dir_name(self) -> str:
        """Return the bin directory name for the platform."""
        return "bin"

    @property
    def host_python_symlink_suffix(self) -> str:
        """Get the host python symlink name."""
        return "-host-link"

    @property
    def host_python_wrapper(self) -> str:
        """Get the content of the host python wrapper script.
        This script allows invoking the host python with the correct PYTHONHOME.
        """
        base_executable = getattr(sys, "_base_executable", sys.executable)
        pythonhome = Path(base_executable).parents[1]
        # These paths can contain spaces, such as noticed via uv-managed
        # Python on macOS which lives under "Application Support". See:
        # https://github.com/pyodide/pyodide-build/issues/399
        return dedent(
            f"""\
            #!/bin/sh
            exec env PYTHONHOME={shlex.quote(str(pythonhome))} {shlex.quote(str(self.host_python_symlink_path))} "$@"
            """
        )

    @property
    def host_pip_wrapper(self) -> str:
        # Other than the shebang and the monkey patch, this is exactly what
        # normal pip looks like.
        return dedent(
            f"""
            #!/usr/bin/env bash
            {shlex.quote(str(self.host_python_path))} -s {shlex.quote(str(self.pip_wrapper_path))} "$@"
            """
        )

    @property
    def host_sitepackages_path(self) -> Path:
        """
        Path to the host-specific site-packages directory in the virtualenv
        """
        return self.venv_sitepackages_path

    def _create_python_symlink(self) -> None:
        """Create a symlink to the Pyodide Python interpreter.

        Noop on Unix as virtualenv already does this for us.
        """
        return

    def _create_session(self):
        """Create and return a virtualenv session object."""
        return virtualenv.session_via_cli(self.get_cli_args() + [str(self.dest)])

    def _create_pyodide_script(self) -> None:
        """Write pyodide cli script into the virtualenv bin folder."""

        # Temporarily restore us to the environment that 'pyodide venv' was
        # invoked in
        PATH = os.environ["PATH"]
        PYODIDE_ROOT = os.environ["PYODIDE_ROOT"]

        original_pyodide_cli = shutil.which("pyodide")
        if original_pyodide_cli is None:
            raise RuntimeError("ERROR: pyodide cli not found")

        self.pyodide_cli_path.write_text(
            dedent(
                f"""
                #!/usr/bin/env bash
                PATH={shlex.quote(PATH)}:"$PATH" PYODIDE_ROOT={shlex.quote(PYODIDE_ROOT)} exec {shlex.quote(original_pyodide_cli)} "$@"
                """
            )
        )
        self.pyodide_cli_path.chmod(0o777)


class WindowsPyodideVenv(PyodideVenv):
    """Windows-specific implementation of Pyodide virtual environment creation."""

    @property
    def exe_suffix(self) -> str:
        """Return the executable suffix for the platform."""
        return ".bat"

    @property
    def bin_dir_name(self) -> str:
        """Return the bin directory name."""
        return "Scripts"

    @property
    def host_python_symlink_suffix(self) -> str:
        """Get the host python symlink name."""
        return "-host-link.exe"

    @property
    def script_interpreter_suffix(self) -> str:
        """Get the extension of the interpreter that console scripts invoke.

        The launcher that pip puts in front of a console script can only
        invoke an ``.exe``, so these point at the ``python.exe`` shim rather
        than at ``python.bat``.
        """
        return ".exe"

    @property
    def interpreter_launcher_path(self) -> Path:
        """Get the path to the python.exe shim in the Pyodide dist directory."""
        return pyodide_dist_dir() / "python.exe"

    @property
    def host_python_wrapper(self) -> str:
        """Get the content of the host python wrapper script.

        TODO(check): In windows, it doesn't seem setting PYTHONHOME is required to make it correctly work.
        """
        return dedent(f"""\
            @echo off
            "{self.host_python_symlink_path}" %*
            """)

    @property
    def pip_conf_path(self) -> Path:
        """Get the path to the pip.conf file in the virtualenv."""
        # On windows, pip would look for pip.conf normally, but since we patch pip to mirror unix behavior,
        # we place pip.conf in the venv root directory.
        return self.venv_root / "pip.conf"

    @property
    def host_pip_wrapper(self) -> str:
        return (
            "@echo off\n"
            + f'"{self.host_python_path}" -s '
            + f'"{self.pip_wrapper_path}" %*\n'
        )

    @property
    def host_sitepackages_path(self) -> Path:
        """
        Path to the host-specific site-packages directory in the virtualenv
        """
        return self.venv_root / "Lib" / "site-packages"

    def _create_session(self):
        from .app_data import create_app_data_dir

        cli_args = self.get_cli_args()
        with create_app_data_dir(str(self.interpreter_path)) as app_data_dir:
            cli_args += ["--app-data", app_data_dir]
            session = virtualenv.session_via_cli(cli_args + [str(self.dest)])

        return session

    def _create_python_symlink(self) -> None:
        """Create a symlink to the Pyodide Python interpreter."""

        # the virtualenv does not understand the batch file, so we need to
        # symlink it ourselves
        self.interpreter_symlink_path.unlink(missing_ok=True)
        self.interpreter_symlink_path.symlink_to(self.interpreter_path)

        # Also symlink any other python*.exe files to the interpreter
        other_pythons = self.venv_bin.glob("python*.exe")
        for python_exe in other_pythons:
            python_exe.unlink(missing_ok=True)
            python_bat = python_exe.with_suffix(self.exe_suffix)
            if python_bat != self.interpreter_symlink_path:
                python_bat.unlink(missing_ok=True)
                python_bat.symlink_to(self.interpreter_path)

        self._install_interpreter_launcher()

    def _install_interpreter_launcher(self) -> None:
        """Put the python.exe shim from the Pyodide dist next to python.bat.

        Pyodide ships a small launcher that forwards to python.bat, so that the
        venv has an ``.exe`` interpreter like an ordinary one does. Without it
        the launchers pip generates for console scripts have nothing they can
        invoke, since they refuse any interpreter that is not an ``.exe``.
        """
        if not self.interpreter_launcher_path.exists():
            # Pyodide only started shipping the shim in 0.29.3
            logger.warning(
                "%s not found in the cross-build environment. Console scripts "
                "installed into this virtualenv will not be runnable.",
                self.interpreter_launcher_path.name,
            )
            return

        # Copy rather than symlink. The shim resolves its own path before
        # looking for python.bat beside it, so a symlink would send it to the
        # dist directory and it would lose track of the virtualenv.
        self.script_interpreter_path.unlink(missing_ok=True)
        shutil.copy(self.interpreter_launcher_path, self.script_interpreter_path)

    def _create_pyodide_script(self) -> None:
        """Write pyodide cli script into the virtualenv bin folder."""
        PATH = os.environ["PATH"]
        PYODIDE_ROOT = os.environ["PYODIDE_ROOT"]

        original_pyodide_cli = shutil.which("pyodide")
        if original_pyodide_cli is None:
            raise RuntimeError("ERROR: pyodide cli not found")

        self.pyodide_cli_path.write_text(
            dedent(
                f"""
                @echo off
                set "PATH={PATH};%PATH%"
                set "PYODIDE_ROOT={PYODIDE_ROOT}"
                "{original_pyodide_cli}" %*
                """
            )
        )


def create_pyodide_venv(
    dest: Path, virtualenv_args: list[str] | None = None
) -> PyodideVenv:
    """Create a Pyodide virtualenv and store it into dest"""
    builder = WindowsPyodideVenv if IS_WIN else UnixPyodideVenv
    venv = builder(dest, virtualenv_args)
    venv.create()

    return venv