Support using env() in prelude macro files (pantsbuild#18273)

The recently introduced `env()` function was only implemented for
`BUILD` files (by oversight, hence classified as bug). This fixes the
support for using it also in prelude files (aka macros).

Relates to pantsbuild#18260 
cc @cognifloyd

src/python/pants/engine/internals/build_files.py

@@ -92,6 +92,7 @@ async def evaluate_preludes(
         **parser.builtin_symbols,
     }
     locals: dict[str, Any] = {}
+    env_vars: set[str] = set()
     for file_content in prelude_digest_contents:
         try:
             file_content_str = file_content.content.decode()
@@ -100,14 +101,15 @@ async def evaluate_preludes(
         except Exception as e:
             raise Exception(f"Error parsing prelude file {file_content.path}: {e}")
         error_on_imports(file_content_str, file_content.path)
+        env_vars.update(BUILDFileEnvVarExtractor.get_env_vars(file_content))
     # __builtins__ is a dict, so isn't hashable, and can't be put in a FrozenDict.
     # Fortunately, we don't care about it - preludes should not be able to override builtins, so we just pop it out.
     # TODO: Give a nice error message if a prelude tries to set a expose a non-hashable value.
     locals.pop("__builtins__", None)
     # Ensure preludes can reference each other by populating the shared globals object with references
     # to the other symbols
     globals.update(locals)
-    return BuildFilePreludeSymbols.from_namespace(locals)
+    return BuildFilePreludeSymbols.create(locals, env_vars)
 
 
 @rule
@@ -207,7 +209,7 @@ def visit_Call(self, node: ast.Call):
                 self.env_vars.add(value)
                 return
             else:
-                logging.warning(
+                logger.warning(
                     f"{self.filename}:{arg.lineno}: Only constant string values as variable name to "
                     f"`env()` is currently supported. This `env()` call will always result in "
                     "the default value only."
@@ -218,10 +220,10 @@ def visit_Call(self, node: ast.Call):
 
 
 async def _extract_env_vars(
-    file_content: FileContent, env: CompleteEnvironmentVars
+    file_content: FileContent, extra_env: Sequence[str], env: CompleteEnvironmentVars
 ) -> EnvironmentVars:
     """For BUILD file env vars, we only ever consult the local systems env."""
-    env_vars = BUILDFileEnvVarExtractor.get_env_vars(file_content)
+    env_vars = (*BUILDFileEnvVarExtractor.get_env_vars(file_content), *extra_env)
     return await Get(
         EnvironmentVars,
         {
@@ -300,7 +302,9 @@ async def parse_address_family(
         dependencies_rules_parser_state = None
 
     all_env_vars = [
-        await _extract_env_vars(fc, session_values[CompleteEnvironmentVars])
+        await _extract_env_vars(
+            fc, prelude_symbols.referenced_env_vars, session_values[CompleteEnvironmentVars]
+        )
         for fc in digest_contents
     ]
 

src/python/pants/engine/internals/build_files_test.py

@@ -71,7 +71,7 @@ def test_parse_address_family_empty() -> None:
             ),
             BootstrapStatus(in_progress=False),
             BuildFileOptions(("BUILD",)),
-            BuildFilePreludeSymbols(FrozenDict()),
+            BuildFilePreludeSymbols(FrozenDict(), ()),
             AddressFamilyDir("/dev/null"),
             RegisteredTargetTypes({}),
             UnionMembership({}),
@@ -179,7 +179,7 @@ def test_prelude_check_filepath() -> None:
     )
     with pytest.raises(
         Exception,
-        match="The BUILD file `build_file_dir` may only be used in BUILD files\\. If used",
+        match="The BUILD file symbol `build_file_dir` may only be used in BUILD files\\. If used",
     ):
         run_prelude_parsing_rule(prelude_content)
 
@@ -192,7 +192,7 @@ def test_prelude_check_defaults() -> None:
     )
     with pytest.raises(
         Exception,
-        match="The BUILD file `__defaults__` may only be used in BUILD files\\. If used",
+        match="The BUILD file symbol `__defaults__` may only be used in BUILD files\\. If used",
     ):
         run_prelude_parsing_rule(prelude_content)
 
@@ -205,7 +205,7 @@ def test_prelude_check_env() -> None:
     )
     with pytest.raises(
         Exception,
-        match="The BUILD file `env` may only be used in BUILD files\\.$",
+        match="The BUILD file symbol `env` may only be used in BUILD files\\. If used",
     ):
         run_prelude_parsing_rule(prelude_content)
 
@@ -254,6 +254,17 @@ def macro_func(arg: int) -> str:
     assert {"macro_func"} == set(result.info)
 
 
+def test_prelude_reference_env_vars() -> None:
+    prelude_content = dedent(
+        """
+        def macro():
+            env("MY_ENV")
+        """
+    )
+    result = run_prelude_parsing_rule(prelude_content)
+    assert ("MY_ENV",) == result.referenced_env_vars
+
+
 class ResolveField(StringField):
     alias = "resolve"
 
@@ -581,6 +592,9 @@ def test_macro_undefined_symbol_bootstrap() -> None:
         rules=[QueryRule(AddressFamily, [AddressFamilyDir])],
         is_bootstrap=True,
     )
+    rule_runner.set_options(
+        args=("--build-file-prelude-globs=prelude.py",),
+    )
     rule_runner.write_files(
         {
             "prelude.py": dedent(
@@ -649,6 +663,35 @@ def test_build_file_env_vars(target_adaptor_rule_runner: RuleRunner) -> None:
     assert target_adaptor.kwargs["tags"] == ["default", "tag"]
 
 
+def test_prelude_env_vars(target_adaptor_rule_runner: RuleRunner) -> None:
+    target_adaptor_rule_runner.write_files(
+        {
+            "prelude.py": dedent(
+                """
+                def macro_val():
+                    return env("MACRO_ENV")
+                """
+            ),
+            "BUILD": dedent(
+                """
+                mock_tgt(
+                  description=macro_val(),
+                )
+                """
+            ),
+        },
+    )
+    target_adaptor_rule_runner.set_options(
+        args=("--build-file-prelude-globs=prelude.py",),
+        env={"MACRO_ENV": "from env"},
+    )
+    target_adaptor = target_adaptor_rule_runner.request(
+        TargetAdaptor,
+        [TargetAdaptorRequest(Address(""), description_of_origin="tests")],
+    )
+    assert target_adaptor.kwargs["description"] == "from env"
+
+
 def test_invalid_build_file_env_vars(caplog, target_adaptor_rule_runner: RuleRunner) -> None:
     target_adaptor_rule_runner.write_files(
         {

src/python/pants/engine/internals/mapper_test.py

@@ -42,7 +42,7 @@ def parse_address_map(build_file: str, *, ignore_unrecognized_symbols: bool = Fa
         path,
         build_file,
         parser,
-        BuildFilePreludeSymbols(FrozenDict()),
+        BuildFilePreludeSymbols(FrozenDict(), ()),
         EnvironmentVars({}),
         False,
         BuildFileDefaultsParserState.create(

src/python/pants/engine/internals/parser.py

@@ -11,7 +11,7 @@
 from difflib import get_close_matches
 from io import StringIO
 from pathlib import PurePath
-from typing import Any, Callable, Mapping, TypeVar
+from typing import Any, Callable, Iterable, Mapping, TypeVar
 
 from pants.base.deprecated import warn_or_error
 from pants.base.exceptions import MappingError
@@ -34,17 +34,18 @@
 @dataclass(frozen=True)
 class BuildFilePreludeSymbols:
     info: FrozenDict[str, BuildFileSymbolInfo]
+    referenced_env_vars: tuple[str, ...]
 
     @memoized_property
     def symbols(self) -> FrozenDict[str, Any]:
         return FrozenDict({name: symbol.value for name, symbol in self.info.items()})
 
     @classmethod
-    def from_namespace(cls, ns: Mapping[str, Any]) -> BuildFilePreludeSymbols:
+    def create(cls, ns: Mapping[str, Any], env_vars: Iterable[str]) -> BuildFilePreludeSymbols:
         info = {}
         for name, symb in ns.items():
             info[name] = BuildFileSymbolInfo(name, symb)
-        return cls(info=FrozenDict(info))
+        return cls(info=FrozenDict(info), referenced_env_vars=tuple(sorted(env_vars)))
 
 
 @dataclass(frozen=True)
@@ -119,21 +120,21 @@ def _prelude_check(self, name: str, value: T | None, closure_supported: bool = T
         raise NameError(
             softwrap(
                 f"""
-                The BUILD file {name} may only be used in BUILD files.{note}
+                The BUILD file symbol `{name}` may only be used in BUILD files.{note}
                 """
             )
         )
 
     def filepath(self) -> str:
-        return self._prelude_check("`build_file_dir`", self._filepath)
+        return self._prelude_check("build_file_dir", self._filepath)
 
     @property
     def defaults(self) -> BuildFileDefaultsParserState:
-        return self._prelude_check("`__defaults__`", self._defaults)
+        return self._prelude_check("__defaults__", self._defaults)
 
     @property
     def env_vars(self) -> EnvironmentVars:
-        return self._prelude_check("`env`", self._env_vars, closure_supported=False)
+        return self._prelude_check("env", self._env_vars)
 
     @property
     def is_bootstrap(self) -> bool:

src/python/pants/engine/internals/parser_test.py

@@ -41,7 +41,7 @@ def test_imports_banned(defaults_parser_state: BuildFileDefaultsParserState) ->
         parser.parse(
             "dir/BUILD",
             "\nx = 'hello'\n\nimport os\n",
-            BuildFilePreludeSymbols(FrozenDict()),
+            BuildFilePreludeSymbols(FrozenDict(), ()),
             EnvironmentVars({}),
             False,
             defaults_parser_state,
@@ -67,7 +67,7 @@ def perform_test(extra_targets: list[str], dym: str) -> None:
             object_aliases=build_file_aliases,
             ignore_unrecognized_symbols=False,
         )
-        prelude_symbols = BuildFilePreludeSymbols.from_namespace(FrozenDict({"prelude": 0}))
+        prelude_symbols = BuildFilePreludeSymbols.create({"prelude": 0}, ())
         fmt_extra_sym = str(extra_targets)[1:-1] + (", ") if len(extra_targets) != 0 else ""
         with pytest.raises(ParseError) as exc:
             parser.parse(