Support adding custom env vars (#504)

* Support adding custom env vars

* Review

* Show ruff format diff on CI failure

* Ruff

* add ruff lint for subprocess calls

* add CHANGELOG entry

* allow dict as env in subprocess wrappers

---------

Co-authored-by: David Hewitt <mail@davidhewitt.dev>

Author: konstin

CHANGELOG.md

@@ -4,6 +4,7 @@
 ### Packaging
 - Drop support for Python 3.8. [#479](https://github.com/PyO3/setuptools-rust/pull/479)
 - Support free-threaded Python. [#502](https://github.com/PyO3/setuptools-rust/pull/502)
+- Support adding custom env vars. [#504](https://github.com/PyO3/setuptools-rust/pull/504)
 
 ## 1.10.2 (2024-10-02)
 ### Fixed

emscripten/.ruff.toml

@@ -0,0 +1,2 @@
+[lint]
+extend-ignore = ["TID251"]

noxfile.py

@@ -139,7 +139,7 @@ def test_crossenv(session: nox.Session):
 @nox.session()
 def ruff(session: nox.Session):
     session.install("ruff")
-    session.run("ruff", "format", "--check", ".")
+    session.run("ruff", "format", "--diff", ".")
     session.run("ruff", "check", ".")
 
 

pyproject.toml

@@ -48,9 +48,16 @@ Changelog = "https://github.com/PyO3/setuptools-rust/blob/main/CHANGELOG.md"
 requires = ["setuptools>=62.4", "setuptools_scm"]
 build-backend = "setuptools.build_meta"
 
+[tool.ruff.lint]
+extend-select = ["TID251"]
+
 [tool.ruff.lint.extend-per-file-ignores]
 "__init__.py" = ["F403"]
 
+[tool.ruff.lint.flake8-tidy-imports.banned-api]
+"subprocess.run".msg = "Use `_utils.run_subprocess` to ensure `env` is passed"
+"subprocess.check_output".msg = "Use `_utils.check_subprocess_output` to ensure `env` is passed"
+
 [tool.pytest.ini_options]
 minversion = "6.0"
 addopts = "--doctest-modules"

setuptools_rust/_utils.py

@@ -1,4 +1,48 @@
 import subprocess
+from typing import Any, Optional, Union, cast
+
+
+class Env:
+    """Allow using ``functools.lru_cache`` with an environment variable dictionary.
+
+    Dictionaries are unhashable, but ``functools.lru_cache`` needs all parameters to
+    be hashable, which we solve which a custom ``__hash__``."""
+
+    env: Optional[dict[str, str]]
+
+    def __init__(self, env: Optional[dict[str, str]]):
+        self.env = env
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Env):
+            return False
+        return self.env == other.env
+
+    def __hash__(self) -> int:
+        if self.env is not None:
+            return hash(tuple(sorted(self.env.items())))
+        else:
+            return hash(None)
+
+
+def run_subprocess(
+    *args: Any, env: Union[Env, dict[str, str], None], **kwargs: Any
+) -> subprocess.CompletedProcess:
+    """Wrapper around subprocess.run that requires a decision to pass env."""
+    if isinstance(env, Env):
+        env = env.env
+    kwargs["env"] = env
+    return subprocess.run(*args, **kwargs)  # noqa: TID251 # this is a wrapper to implement the rule
+
+
+def check_subprocess_output(
+    *args: Any, env: Union[Env, dict[str, str], None], **kwargs: Any
+) -> str:
+    """Wrapper around subprocess.run that requires a decision to pass env."""
+    if isinstance(env, Env):
+        env = env.env
+    kwargs["env"] = env
+    return cast(str, subprocess.check_output(*args, **kwargs))  # noqa: TID251 # this is a wrapper to implement the rule
 
 
 def format_called_process_error(

setuptools_rust/build.py

@@ -24,7 +24,7 @@
 from setuptools.command.build_ext import get_abi3_suffix
 from setuptools.command.install_scripts import install_scripts as CommandInstallScripts
 
-from ._utils import format_called_process_error
+from ._utils import check_subprocess_output, format_called_process_error, Env
 from .command import RustCommand
 from .extension import Binding, RustBin, RustExtension, Strip
 from .rustc_info import (
@@ -45,8 +45,8 @@
         from setuptools import Command as CommandBdistWheel  # type: ignore[assignment]
 
 
-def _check_cargo_supports_crate_type_option() -> bool:
-    version = get_rust_version()
+def _check_cargo_supports_crate_type_option(env: Optional[Env]) -> bool:
+    version = get_rust_version(env)
 
     if version is None:
         return False
@@ -144,10 +144,10 @@ def run_for_extension(self, ext: RustExtension) -> None:
     def build_extension(
         self, ext: RustExtension, forced_target_triple: Optional[str] = None
     ) -> List["_BuiltModule"]:
-        target_triple = self._detect_rust_target(forced_target_triple)
-        rustc_cfgs = get_rustc_cfgs(target_triple)
+        target_triple = self._detect_rust_target(forced_target_triple, ext.env)
+        rustc_cfgs = get_rustc_cfgs(target_triple, ext.env)
 
-        env = _prepare_build_environment()
+        env = _prepare_build_environment(ext.env)
 
         if not os.path.exists(ext.path):
             raise FileError(
@@ -156,7 +156,7 @@ def build_extension(
 
         quiet = self.qbuild or ext.quiet
         debug = self._is_debug_build(ext)
-        use_cargo_crate_type = _check_cargo_supports_crate_type_option()
+        use_cargo_crate_type = _check_cargo_supports_crate_type_option(ext.env)
 
         package_id = ext.metadata(quiet=quiet)["resolve"]["root"]
         if package_id is None:
@@ -252,7 +252,7 @@ def build_extension(
             # If quiet, capture all output and only show it in the exception
             # If not quiet, forward all cargo output to stderr
             stderr = subprocess.PIPE if quiet else None
-            cargo_messages = subprocess.check_output(
+            cargo_messages = check_subprocess_output(
                 command,
                 env=env,
                 stderr=stderr,
@@ -417,7 +417,7 @@ def install_extension(
                     args.insert(0, "strip")
                     args.append(ext_path)
                     try:
-                        subprocess.check_output(args)
+                        check_subprocess_output(args, env=None)
                     except subprocess.CalledProcessError:
                         pass
 
@@ -477,7 +477,7 @@ def _py_limited_api(self) -> _PyLimitedApi:
             return cast(_PyLimitedApi, bdist_wheel.py_limited_api)
 
     def _detect_rust_target(
-        self, forced_target_triple: Optional[str] = None
+        self, forced_target_triple: Optional[str], env: Env
     ) -> Optional[str]:
         assert self.plat_name is not None
         if forced_target_triple is not None:
@@ -486,14 +486,14 @@ def _detect_rust_target(
             return forced_target_triple
 
         # Determine local rust target which needs to be "forced" if necessary
-        local_rust_target = _adjusted_local_rust_target(self.plat_name)
+        local_rust_target = _adjusted_local_rust_target(self.plat_name, env)
 
         # Match cargo's behaviour of not using an explicit target if the
         # target we're compiling for is the host
         if (
             local_rust_target is not None
             # check for None first to avoid calling to rustc if not needed
-            and local_rust_target != get_rust_host()
+            and local_rust_target != get_rust_host(env)
         ):
             return local_rust_target
 
@@ -566,7 +566,7 @@ def create_universal2_binary(output_path: str, input_paths: List[str]) -> None:
     # Try lipo first
     command = ["lipo", "-create", "-output", output_path, *input_paths]
     try:
-        subprocess.check_output(command, text=True)
+        check_subprocess_output(command, env=None, text=True)
     except subprocess.CalledProcessError as e:
         output = e.output
         raise CompileError("lipo failed with code: %d\n%s" % (e.returncode, output))
@@ -609,7 +609,7 @@ def _replace_vendor_with_unknown(target: str) -> Optional[str]:
     return "-".join(components)
 
 
-def _prepare_build_environment() -> Dict[str, str]:
+def _prepare_build_environment(env: Env) -> Dict[str, str]:
     """Prepares environment variables to use when executing cargo build."""
 
     base_executable = None
@@ -625,20 +625,18 @@ def _prepare_build_environment() -> Dict[str, str]:
     # executing python interpreter.
     bindir = os.path.dirname(executable)
 
-    env = os.environ.copy()
-    env.update(
+    env_vars = (env.env or os.environ).copy()
+    env_vars.update(
         {
             # disables rust's pkg-config seeking for specified packages,
             # which causes pythonXX-sys to fall back to detecting the
             # interpreter from the path.
-            "PATH": os.path.join(bindir, os.environ.get("PATH", "")),
-            "PYTHON_SYS_EXECUTABLE": os.environ.get(
-                "PYTHON_SYS_EXECUTABLE", executable
-            ),
-            "PYO3_PYTHON": os.environ.get("PYO3_PYTHON", executable),
+            "PATH": os.path.join(bindir, env_vars.get("PATH", "")),
+            "PYTHON_SYS_EXECUTABLE": env_vars.get("PYTHON_SYS_EXECUTABLE", executable),
+            "PYO3_PYTHON": env_vars.get("PYO3_PYTHON", executable),
         }
     )
-    return env
+    return env_vars
 
 
 def _is_py_limited_api(
@@ -692,19 +690,19 @@ def _binding_features(
 _PyLimitedApi = Literal["cp37", "cp38", "cp39", "cp310", "cp311", "cp312", True, False]
 
 
-def _adjusted_local_rust_target(plat_name: str) -> Optional[str]:
+def _adjusted_local_rust_target(plat_name: str, env: Env) -> Optional[str]:
     """Returns the local rust target for the given `plat_name`, if it is
     necessary to 'force' a specific target for correctness."""
 
     # If we are on a 64-bit machine, but running a 32-bit Python, then
     # we'll target a 32-bit Rust build.
     if plat_name == "win32":
-        if get_rustc_cfgs(None).get("target_env") == "gnu":
+        if get_rustc_cfgs(None, env).get("target_env") == "gnu":
             return "i686-pc-windows-gnu"
         else:
             return "i686-pc-windows-msvc"
     elif plat_name == "win-amd64":
-        if get_rustc_cfgs(None).get("target_env") == "gnu":
+        if get_rustc_cfgs(None, env).get("target_env") == "gnu":
             return "x86_64-pc-windows-gnu"
         else:
             return "x86_64-pc-windows-msvc"

setuptools_rust/clean.py

@@ -1,6 +1,7 @@
-import subprocess
 import sys
 
+from setuptools_rust._utils import check_subprocess_output
+
 from .command import RustCommand
 from .extension import RustExtension
 
@@ -25,6 +26,6 @@ def run_for_extension(self, ext: RustExtension) -> None:
 
         # Execute cargo command
         try:
-            subprocess.check_output(args)
+            check_subprocess_output(args, env=ext.env)
         except Exception:
             pass

setuptools_rust/command.py

@@ -52,8 +52,16 @@ def run(self) -> None:
             return
 
         all_optional = all(ext.optional for ext in self.extensions)
+        # Use the environment of the first non-optional extension, or the first optional
+        # extension if there is no non-optional extension.
+        env = None
+        for ext in self.extensions:
+            if ext.env:
+                env = ext.env
+                if not ext.optional:
+                    break
         try:
-            version = get_rust_version()
+            version = get_rust_version(env)
             if version is None:
                 min_version = max(  # type: ignore[type-var]
                     filter(

setuptools_rust/extension.py

@@ -24,7 +24,7 @@
 if TYPE_CHECKING:
     from semantic_version import SimpleSpec
 
-from ._utils import format_called_process_error
+from ._utils import check_subprocess_output, format_called_process_error, Env
 
 
 class Binding(IntEnum):
@@ -112,6 +112,9 @@ class RustExtension:
             abort the build process, and instead simply not install the failing
             extension.
         py_limited_api: Deprecated.
+        env: Environment variables to use when calling cargo or rustc (``env=``
+            in ``subprocess.Popen``). setuptools-rust may add additional
+            variables or modify ``PATH``.
     """
 
     def __init__(
@@ -131,6 +134,7 @@ def __init__(
         native: bool = False,
         optional: bool = False,
         py_limited_api: Literal["auto", True, False] = "auto",
+        env: Optional[Dict[str, str]] = None,
     ):
         if isinstance(target, dict):
             name = "; ".join("%s=%s" % (key, val) for key, val in target.items())
@@ -153,6 +157,7 @@ def __init__(
         self.script = script
         self.optional = optional
         self.py_limited_api = py_limited_api
+        self.env = Env(env)
 
         if native:
             warnings.warn(
@@ -260,8 +265,8 @@ def _metadata(self, cargo: str, quiet: bool) -> "CargoMetadata":
             # If quiet, capture stderr and only show it on exceptions
             # If not quiet, let stderr be inherited
             stderr = subprocess.PIPE if quiet else None
-            payload = subprocess.check_output(
-                metadata_command, stderr=stderr, encoding="latin-1"
+            payload = check_subprocess_output(
+                metadata_command, stderr=stderr, encoding="latin-1", env=self.env.env
             )
         except subprocess.CalledProcessError as e:
             raise SetupError(format_called_process_error(e))
@@ -319,6 +324,7 @@ def __init__(
         debug: Optional[bool] = None,
         strip: Strip = Strip.No,
         optional: bool = False,
+        env: Optional[dict[str, str]] = None,
     ):
         super().__init__(
             target=target,
@@ -333,6 +339,7 @@ def __init__(
             optional=optional,
             strip=strip,
             py_limited_api=False,
+            env=env,
         )
 
     def entry_points(self) -> List[str]: