Optionally inject environment variables into EnvironmentAware

src/python/pants/backend/python/subsystems/python_native_code.py

@@ -3,21 +3,21 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import Dict, Sequence
+from typing import Sequence
 
-from pants.engine.env_vars import EnvironmentVars, EnvironmentVarsRequest
-from pants.engine.rules import Get, collect_rules, rule
 from pants.option.option_types import StrListOption
 from pants.option.subsystem import Subsystem
 from pants.util.strutil import safe_shlex_join, safe_shlex_split
 
 
 class PythonNativeCodeSubsystem(Subsystem):
+
     options_scope = "python-native-code"
     help = "Options for building native code using Python, e.g. when resolving distributions."
 
     class EnvironmentAware(Subsystem.EnvironmentAware):
+        depends_on_env_vars = ("CPPFLAGS", "LDFLAGS")
+
         # TODO(#7735): move the --cpp-flags and --ld-flags to a general subprocess support subsystem.
         _cpp_flags = StrListOption(
             default=["<CPPFLAGS>"],
@@ -38,40 +38,16 @@ class EnvironmentAware(Subsystem.EnvironmentAware):
             advanced=True,
         )
 
-
-@dataclass(frozen=True)
-class PythonNativeCodeEnvironment:
-
-    cpp_flags: tuple[str, ...]
-    ld_flags: tuple[str, ...]
-
-    @property
-    def environment_dict(self) -> Dict[str, str]:
-        return {
-            "CPPFLAGS": safe_shlex_join(self.cpp_flags),
-            "LDFLAGS": safe_shlex_join(self.ld_flags),
-        }
-
-
-@rule
-async def resolve_python_native_code_environment(
-    env_aware: PythonNativeCodeSubsystem.EnvironmentAware,
-) -> PythonNativeCodeEnvironment:
-
-    env_vars = await Get(EnvironmentVars, EnvironmentVarsRequest(("CPPFLAGS", "LDFLAGS")))
-
-    def iter_values(env_var: str, values: Sequence[str]):
-        for value in values:
-            if value == f"<{env_var}>":
-                yield from safe_shlex_split(env_vars.get(env_var, ""))
-            else:
-                yield value
-
-    return PythonNativeCodeEnvironment(
-        cpp_flags=tuple(iter_values("CPPFLAGS", env_aware._cpp_flags)),
-        ld_flags=tuple(iter_values("LDFLAGS", env_aware._ld_flags)),
-    )
-
-
-def rules():
-    return [*collect_rules()]
+        @property
+        def subprocess_env_vars(self) -> dict[str, str]:
+            return {
+                "CPPFLAGS": safe_shlex_join(self._iter_values("CPPFLAGS", self._cpp_flags)),
+                "LDFLAGS": safe_shlex_join(self._iter_values("LDFLAGS", self._ld_flags)),
+            }
+
+        def _iter_values(self, env_var: str, values: Sequence[str]):
+            for value in values:
+                if value == f"<{env_var}>":
+                    yield from safe_shlex_split(self.env_vars.get(env_var, ""))
+                else:
+                    yield value

src/python/pants/backend/python/util_rules/pex_cli.py

@@ -9,8 +9,7 @@
 from pathlib import Path
 from typing import Iterable, List, Mapping, Optional, Tuple
 
-from pants.backend.python.subsystems import python_native_code
-from pants.backend.python.subsystems.python_native_code import PythonNativeCodeEnvironment
+from pants.backend.python.subsystems.python_native_code import PythonNativeCodeSubsystem
 from pants.backend.python.util_rules import pex_environment
 from pants.backend.python.util_rules.pex_environment import (
     PexEnvironment,
@@ -125,7 +124,7 @@ async def setup_pex_cli_process(
     request: PexCliProcess,
     pex_pex: PexPEX,
     pex_env: PexEnvironment,
-    python_native_code_environment: PythonNativeCodeEnvironment,
+    python_native_code: PythonNativeCodeSubsystem.EnvironmentAware,
     global_options: GlobalOptions,
     pex_subsystem: PexSubsystem,
 ) -> Process:
@@ -191,7 +190,7 @@ async def setup_pex_cli_process(
     normalized_argv = complete_pex_env.create_argv(pex_pex.exe, *args, python=request.python)
     env = {
         **complete_pex_env.environment_dict(python_configured=request.python is not None),
-        **python_native_code_environment.environment_dict,
+        **python_native_code.subprocess_env_vars,
         **(request.extra_env or {}),
         # If a subcommand is used, we need to use the `pex3` console script.
         **({"PEX_SCRIPT": "pex3"} if request.subcommand else {}),
@@ -216,5 +215,4 @@ def rules():
         *collect_rules(),
         *external_tool.rules(),
         *pex_environment.rules(),
-        *python_native_code.rules(),
     ]

src/python/pants/backend/shell/shell_command.py

@@ -133,8 +133,7 @@ async def prepare_shell_command_process(
                 f"Must provide any `tools` used by the `{shell_command.alias}` {shell_command.address}."
             )
 
-        env = await Get(EnvironmentVars, EnvironmentVarsRequest(["PATH"]))
-        search_path = shell_setup.executable_search_path(env)
+        search_path = shell_setup.executable_search_path
         tool_requests = [
             BinaryPathRequest(
                 binary_name=tool,

src/python/pants/backend/shell/shell_setup.py

@@ -5,10 +5,9 @@
 
 import os
 
-from pants.engine.env_vars import EnvironmentVars
 from pants.option.option_types import BoolOption, StrListOption
 from pants.option.subsystem import Subsystem
-from pants.util.memo import memoized_method
+from pants.util.memo import memoized_property
 from pants.util.ordered_set import OrderedSet
 from pants.util.strutil import softwrap
 
@@ -32,7 +31,9 @@ class ShellSetup(Subsystem):
         advanced=True,
     )
 
-    class EnvironmentAware:
+    class EnvironmentAware(Subsystem.EnvironmentAware):
+        depends_on_env_vars = ("PATH",)
+
         _executable_search_path = StrListOption(
             default=["<PATH>"],
             help=softwrap(
@@ -47,12 +48,12 @@ class EnvironmentAware:
             metavar="<binary-paths>",
         )
 
-        @memoized_method
-        def executable_search_path(self, env: EnvironmentVars) -> tuple[str, ...]:
+        @memoized_property
+        def executable_search_path(self) -> tuple[str, ...]:
             def iter_path_entries():
                 for entry in self._executable_search_path:
                     if entry == "<PATH>":
-                        path = env.get("PATH")
+                        path = self.env_vars.get("PATH")
                         if path:
                             yield from path.split(os.pathsep)
                     else:

src/python/pants/backend/shell/shunit2_test_runner.py

@@ -34,7 +34,6 @@
     BinaryPaths,
 )
 from pants.engine.addresses import Address
-from pants.engine.env_vars import EnvironmentVars, EnvironmentVarsRequest
 from pants.engine.fs import (
     CreateDigest,
     Digest,
@@ -141,10 +140,9 @@ async def determine_shunit2_shell(
             )
         tgt_shell = parse_result
 
-    env = await Get(EnvironmentVars, EnvironmentVarsRequest(["PATH"]))
     path_request = BinaryPathRequest(
         binary_name=tgt_shell.name,
-        search_path=shell_setup.executable_search_path(env),
+        search_path=shell_setup.executable_search_path,
         test=tgt_shell.binary_path_test,
     )
     paths = await Get(BinaryPaths, BinaryPathRequest, path_request)
@@ -168,14 +166,13 @@ async def setup_shunit2_for_target(
         "https://raw.githubusercontent.com/kward/shunit2/b9102bb763cc603b3115ed30a5648bf950548097/shunit2",
         FileDigest("1f11477b7948150d1ca50cdd41d89be4ed2acd137e26d2e0fe23966d0e272cc5", 40987),
     )
-    shunit2_script, transitive_targets, built_package_dependencies, env = await MultiGet(
+    shunit2_script, transitive_targets, built_package_dependencies = await MultiGet(
         Get(Digest, DownloadFile, shunit2_download_file),
         Get(TransitiveTargets, TransitiveTargetsRequest([request.field_set.address])),
         Get(
             BuiltPackageDependencies,
             BuildPackageDependenciesRequest(request.field_set.runtime_package_dependencies),
         ),
-        Get(EnvironmentVars, EnvironmentVarsRequest(["PATH"])),
     )
 
     dependencies_source_files_request = Get(
@@ -219,7 +216,7 @@ async def setup_shunit2_for_target(
     )
 
     env_dict = {
-        "PATH": create_path_env_var(shell_setup.executable_search_path(env)),
+        "PATH": create_path_env_var(shell_setup.executable_search_path),
         "SHUNIT_COLOR": "always" if global_options.colors else "none",
         **test_extra_env.env,
     }

src/python/pants/option/subsystem.py

@@ -11,6 +11,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Iterable, TypeVar, cast
 
 from pants import ox
+from pants.engine.env_vars import EnvironmentVars, EnvironmentVarsRequest
 from pants.engine.internals.selectors import AwaitableConstraints, Get
 from pants.option.errors import OptionsError
 from pants.option.option_types import OptionsInfo, collect_options_info
@@ -87,8 +88,11 @@ class EnvironmentAware(metaclass=ABCMeta):
         """
 
         subsystem: ClassVar[type[Subsystem]]
+        depends_on_env_vars: ClassVar[tuple[str, ...]] = ()
+
         options: OptionValueContainer
         env_tgt: EnvironmentTarget
+        env_vars: EnvironmentVars = EnvironmentVars()
 
         def __getattribute__(self, __name: str) -> Any:
             from pants.core.util_rules.environments import resolve_environment_sensitive_option
@@ -191,7 +195,13 @@ async def inner(*a, **k):
             output_type=cls.EnvironmentAware,
             input_selectors=(cls, EnvironmentTarget),
             func=inner,
-            input_gets=(),
+            input_gets=(
+                AwaitableConstraints(
+                    output_type=EnvironmentVars,
+                    input_types=(EnvironmentVarsRequest,),
+                    is_effect=False,
+                ),
+            ),
             canonical_name=name,
         )
 
@@ -266,4 +276,7 @@ async def _construct_env_aware(
     t.options = subsystem_instance.options
     t.env_tgt = env_tgt
 
+    if t.depends_on_env_vars:
+        t.env_vars = await Get(EnvironmentVars, EnvironmentVarsRequest(t.depends_on_env_vars))
+
     return t