Allow passing additional environment variables to venv (#19)

This allows to specify environment variables that are added to the virtual python environment when build commands are run.
This feature is now available in `VirtualPythonEnvironment` and all derivatives.

* README.md : Add example using the `env` keyword.
* sphinx_polyversion/pyenv.py (VirtualPythonEnvironment): Add `env` keyword and implement in the new `apply_overrides` method called by `run`.
* sphinx_polyversion/pyenv.py : Add the `env` keyword to the subclasses `Poetry` and `Pip` as well.
* tests/test_pyenv.py : Test the new behaviour.

---------

Co-authored-by: real-yfprojects <[email protected]>

Author: vpratz

README.md

@@ -185,7 +185,7 @@ BUILDER = {
 ENVIRONMENT = {
     None: Poetry.factory(args="--sync".split()),  # first version
     "v1.5.7": Poetry.factory(args="--only sphinx --sync".split()),
-    "v1.8.2": Poetry.factory(args="--only dev --sync".split()),
+    "v1.8.2": Poetry.factory(args="--only dev --sync".split(), env={"MY_VAR": "value"}),
     "v3.0.0": Pip.factory(venv=Path(".venv"), args="-e . -r requirements.txt".split()),
 }
 

sphinx_polyversion/pyvenv.py

@@ -99,6 +99,9 @@ class VirtualPythonEnvironment(Environment):
         The path of the python venv.
     creator : Callable[[Path], Any], optional
         A callable for creating the venv, by default None
+    env  : dict[str, str], optional
+        A dictionary of environment variables which are overridden in the
+        virtual environment, by default None
 
     Attributes
     ----------
@@ -108,6 +111,8 @@ class VirtualPythonEnvironment(Environment):
         The name of the environment.
     venv : Path
         The path of the python venv.
+    env  : dict
+        The user-specified environment variables for the virtual environment.
 
     """
 
@@ -118,6 +123,7 @@ def __init__(
         venv: str | Path,
         *,
         creator: Callable[[Path], Any] | None = None,
+        env: dict[str, str] | None = None,
     ):
         """
         Environment for building inside a python virtual environment.
@@ -132,11 +138,15 @@ def __init__(
             The path of the python venv.
         creator : Callable[[Path], Any], optional
             A callable for creating the venv, by default None
+        env  : dict[str, str], optional
+            A dictionary of environment variables which are forwarded to the
+            virtual environment, by default None
 
         """
         super().__init__(path, name)
         self.venv = Path(venv).resolve()
         self._creator = creator
+        self.env = env or {}
 
     async def create_venv(self) -> None:
         """
@@ -191,6 +201,45 @@ def activate(self, env: dict[str, str]) -> dict[str, str]:
         env["PATH"] = str(self.venv / "bin") + ":" + env["PATH"]
         return env
 
+    def apply_overrides(self, env: dict[str, str]) -> dict[str, str]:
+        """
+        Prepare the environment for the build.
+
+        This method is used to modify the environment before running a
+        build command. It :py:meth:`activates <activate>` the python venv
+        and overrides those environment variables that were passed to the
+        :py:class:`constructor <VirtualPythonEnvironment>`.
+        `PATH` is never replaced but extended instead.
+
+        .. warning:: This modifies the given dictionary in-place.
+
+        Parameters
+        ----------
+        env : dict[str, str]
+            The environment to modify.
+
+        Returns
+        -------
+        dict[str, str]
+            The updated environment.
+
+        """
+        # add user-supplied values to env
+        for key, value in self.env.items():
+            if key == "PATH":
+                # extend PATH instead of replacing
+                env["PATH"] = value + ":" + env["PATH"]
+                continue
+            if key in env:
+                logger.info(
+                    "Overwriting environment variable %s=%s with user-specified value '%s'.",
+                    key,
+                    env[key],
+                    value,
+                )
+            env[key] = value
+        return env
+
     async def run(
         self, *cmd: str, **kwargs: Any
     ) -> Tuple[str | bytes | None, str | bytes | None, int]:
@@ -199,8 +248,9 @@ async def run(
 
         This implementation passes the arguments to
         :func:`asyncio.create_subprocess_exec`. But alters `env` to
-        activate the correct python venv. If a python venv is already activated
-        this activation is overridden.
+        :py:meth:`activates <activate>` the correct python
+        and overrides the use-specified vars using :py:meth:`prepare_env`.
+        If a python venv is already activated this activation is overridden.
 
         Returns
         -------
@@ -213,7 +263,9 @@ async def run(
 
         """
         # activate venv
-        kwargs["env"] = self.activate(kwargs.get("env", os.environ).copy())
+        kwargs["env"] = self.activate(
+            self.apply_overrides(kwargs.get("env", os.environ).copy())
+        )
         return await super().run(*cmd, **kwargs)
 
 
@@ -232,10 +284,20 @@ class Poetry(VirtualPythonEnvironment):
         The name of the environment (usually the name of the revision).
     args : Iterable[str]
         The cmd arguments to pass to `poetry install`.
+    env  : dict[str, str], optional
+        A dictionary of environment variables which are overidden in the
+        virtual environment, by default None
 
     """
 
-    def __init__(self, path: Path, name: str, *, args: Iterable[str]):
+    def __init__(
+        self,
+        path: Path,
+        name: str,
+        *,
+        args: Iterable[str],
+        env: dict[str, str] | None = None,
+    ):
         """
         Build Environment for isolated builds using poetry.
 
@@ -247,12 +309,16 @@ def __init__(self, path: Path, name: str, *, args: Iterable[str]):
             The name of the environment (usually the name of the revision).
         args : Iterable[str]
             The cmd arguments to pass to `poetry install`.
+        env  : dict[str, str], optional
+            A dictionary of environment variables which are forwarded to the
+            virtual environment, by default None
 
         """
         super().__init__(
             path,
             name,
             path / ".venv",  # placeholder, determined later
+            env=env,
         )
         self.args = args
 
@@ -273,6 +339,8 @@ async def __aenter__(self) -> Self:
         cmd += self.args
 
         env = os.environ.copy()
+        self.apply_overrides(env)
+
         env.pop("VIRTUAL_ENV", None)  # unset poetry env
         env["POETRY_VIRTUALENVS_IN_PROJECT"] = "False"
         venv_path = self.path / ".venv"
@@ -344,6 +412,9 @@ class Pip(VirtualPythonEnvironment):
         The cmd arguments to pass to `pip install`.
     creator : Callable[[Path], Any] | None, optional
         A callable for creating the venv, by default None
+    env  : dict[str, str], optional
+        A dictionary of environment variables which are overridden in the
+        virtual environment, by default None
 
     """
 
@@ -355,6 +426,7 @@ def __init__(
         *,
         args: Iterable[str],
         creator: Callable[[Path], Any] | None = None,
+        env: dict[str, str] | None = None,
     ):
         """
         Build Environment for using a venv and pip.
@@ -371,9 +443,12 @@ def __init__(
             The cmd arguments to pass to `pip install`.
         creator : Callable[[Path], Any], optional
             A callable for creating the venv, by default None
+        env  : dict[str, str], optional
+            A dictionary of environment variables which are overridden in the
+            virtual environment, by default None
 
         """
-        super().__init__(path, name, venv, creator=creator)
+        super().__init__(path, name, venv, creator=creator, env=env)
         self.args = args
 
     async def __aenter__(self) -> Self:
@@ -396,7 +471,7 @@ async def __aenter__(self) -> Self:
         process = await asyncio.create_subprocess_exec(
             *cmd,
             cwd=self.path,
-            env=self.activate(os.environ.copy()),
+            env=self.activate(self.apply_overrides(os.environ.copy())),
             stdout=PIPE,
             stderr=PIPE,
         )

tests/test_pyvenv.py

@@ -106,6 +106,23 @@ async def test_run_with_creator(self, tmp_path: Path):
             assert rc == 0
             assert str(location) == out.strip()
 
+    @pytest.mark.asyncio
+    async def test_run_with_env_variables(self, tmp_path: Path):
+        """Test passing an environment variable to a venv."""
+        location = tmp_path / "venv"
+
+        async with VirtualPythonEnvironment(
+            tmp_path, "main", location, creator=VenvWrapper(), env={"TESTVAR": "1"}
+        ) as env:
+            out, err, rc = await env.run(
+                "python",
+                "-c",
+                "import os; print(os.environ['TESTVAR'])",
+                stdout=asyncio.subprocess.PIPE,
+            )
+            assert rc == 0
+            assert out.strip() == "1"
+
 
 class TestPip:
     """Test the `Pip` class."""
@@ -139,6 +156,38 @@ async def test_install_into_existing_venv(self, tmp_path: Path):
             )
             assert rc == 0
 
+    @pytest.mark.asyncio
+    async def test_run_with_env_variables(self, tmp_path: Path):
+        """Test passing an environment variable to a venv."""
+        location = tmp_path / "venv"
+
+        # create env
+        await VenvWrapper(with_pip=True)(location)
+
+        # test that tomli is not installed
+        proc = await asyncio.create_subprocess_exec(
+            str(location / "bin/python"),
+            "-c",
+            "import tomli",
+            stdout=asyncio.subprocess.PIPE,
+        )
+        rc = await proc.wait()
+        assert rc == 1
+
+        # init env with tomli
+        async with Pip(
+            tmp_path, "main", location, args=["tomli"], env={"TESTVAR": "1"}
+        ) as env:
+            # test that tomli is installed
+            out, err, rc = await env.run(
+                "python",
+                "-c",
+                "import os; print(os.environ['TESTVAR'])",
+                stdout=asyncio.subprocess.PIPE,
+            )
+            assert rc == 0
+            assert out.strip() == "1"
+
 
 class TestPoetry:
     """Test the `Poetry` environment."""
@@ -169,7 +218,7 @@ async def test_simple_project(self, tmp_path: Path):
         )
 
         # create poetry env
-        async with Poetry(tmp_path, "main", args=[]) as env:
+        async with Poetry(tmp_path, "main", args=[], env={"TESTVAR": "1"}) as env:
             # check sourcing works
             out, err, rc = await env.run(
                 "python",
@@ -198,6 +247,16 @@ async def test_simple_project(self, tmp_path: Path):
             )
             assert rc == 0
 
+            # test that custom environment variables are passed correctly
+            out, err, rc = await env.run(
+                "python",
+                "-c",
+                "import os; print(os.environ['TESTVAR'])",
+                stdout=asyncio.subprocess.PIPE,
+            )
+            assert rc == 0
+            assert out.strip() == "1"
+
     @pytest.mark.asyncio
     async def test_simple_project_with_optional_deps(self, tmp_path: Path):
         """Test installing a simple project with poetry."""