Optionally inject environment variables into EnvironmentAware (#17013)

This extends #17009 by making a generic mechanism to request environment variables when `EnvironmentAware` subsystems are constructed.

Generally speaking, this encapsulates both the requesting and consumption of environment variables into the subsystem itself, and this will save subsystem consumers from knowing which environment variables to request in order to call an option post-processing method. All option post-processing methods should safely be able to be converted to properties too. All very exciting.

Fixes #14612

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