From aaa50d658567a0fa913932e7b430d17bbfbd76be Mon Sep 17 00:00:00 2001 From: Andreas Stenius Date: Tue, 29 Nov 2022 09:30:16 -0500 Subject: [PATCH 1/2] [Plugin API] Add type `CurrentExecutingGoal` to allow rules to find out the current goal. --- src/python/pants/bin/local_pants_runner.py | 2 ++ src/python/pants/engine/goal.py | 23 +++++++++++++ src/python/pants/engine/goal_test.py | 33 +++++++++++++++++-- .../pants/engine/internals/scheduler.py | 26 +++++++++------ src/python/pants/init/engine_initializer.py | 6 +++- src/python/pants/testutil/rule_runner.py | 3 +- 6 files changed, 78 insertions(+), 15 deletions(-) diff --git a/src/python/pants/bin/local_pants_runner.py b/src/python/pants/bin/local_pants_runner.py index 7b62f96a864..e037ba74e5f 100644 --- a/src/python/pants/bin/local_pants_runner.py +++ b/src/python/pants/bin/local_pants_runner.py @@ -13,6 +13,7 @@ from pants.build_graph.build_configuration import BuildConfiguration from pants.core.util_rules.environments import determine_bootstrap_environment from pants.engine.env_vars import CompleteEnvironmentVars +from pants.engine.goal import CurrentExecutingGoal from pants.engine.internals import native_engine from pants.engine.internals.native_engine import PySessionCancellationLatch from pants.engine.internals.scheduler import ExecutionError @@ -100,6 +101,7 @@ def _init_graph_session( { OptionsBootstrapper: options_bootstrapper, CompleteEnvironmentVars: env, + CurrentExecutingGoal: CurrentExecutingGoal(), } ), cancellation_latch=cancellation_latch, diff --git a/src/python/pants/engine/goal.py b/src/python/pants/engine/goal.py index 61c8b40ec7c..e6066a80828 100644 --- a/src/python/pants/engine/goal.py +++ b/src/python/pants/engine/goal.py @@ -1,5 +1,6 @@ # Copyright 2019 Pants project contributors (see CONTRIBUTORS.md). # Licensed under the Apache License, Version 2.0 (see LICENSE). +from __future__ import annotations from abc import abstractmethod from contextlib import contextmanager @@ -10,6 +11,7 @@ from typing_extensions import final from pants.base.deprecated import deprecated_conditional +from pants.engine.engine_aware import EngineAwareReturnType from pants.engine.unions import UnionMembership from pants.option.option_types import StrOption from pants.option.scope import ScopeInfo @@ -133,6 +135,27 @@ def name(cls) -> str: return cast(str, cls.subsystem_cls.name) +@dataclass(frozen=True) +class CurrentExecutingGoal(EngineAwareReturnType): + goal: type[Goal] | None = None + + @property + def name(self) -> str | None: + return None if self.goal is None else self.goal.name + + @contextmanager + def _executing(self, goal: type[Goal]) -> Iterator[None]: + # Mutate current goal; we're only frozen to avoid inadvertent tampering with `self.goal`. + object.__setattr__(self, "goal", goal) + try: + yield + finally: + object.__setattr__(self, "goal", None) + + def cacheable(self) -> bool: + return False + + class Outputting: """A mixin for Goal that adds options to support output-related context managers. diff --git a/src/python/pants/engine/goal_test.py b/src/python/pants/engine/goal_test.py index 5f296ed266e..7e1aa19a529 100644 --- a/src/python/pants/engine/goal_test.py +++ b/src/python/pants/engine/goal_test.py @@ -2,11 +2,11 @@ # Licensed under the Apache License, Version 2.0 (see LICENSE). from pants.engine.console import Console -from pants.engine.goal import Goal, GoalSubsystem, LineOriented -from pants.engine.rules import goal_rule +from pants.engine.goal import CurrentExecutingGoal, Goal, GoalSubsystem, LineOriented +from pants.engine.rules import collect_rules, goal_rule from pants.option.scope import ScopeInfo from pants.testutil.option_util import create_goal_subsystem, create_options_bootstrapper -from pants.testutil.rule_runner import mock_console, run_rule_with_mocks +from pants.testutil.rule_runner import RuleRunner, mock_console, run_rule_with_mocks def test_line_oriented_goal() -> None: @@ -43,3 +43,30 @@ class DummyGoal(GoalSubsystem): dummy = create_goal_subsystem(DummyGoal) assert dummy.get_scope_info() == ScopeInfo(scope="dummy", subsystem_cls=DummyGoal, is_goal=True) + + +def test_current_executing_goal() -> None: + class OutputtingGoalOptions(LineOriented, GoalSubsystem): + name = "dummy" + help = "dummy help" + + class OutputtingGoal(Goal): + subsystem_cls = OutputtingGoalOptions + environment_behavior = Goal.EnvironmentBehavior.LOCAL_ONLY + + @goal_rule + def output_rule( + console: Console, + options: OutputtingGoalOptions, + current_executing_goal: CurrentExecutingGoal, + ) -> OutputtingGoal: + with options.output(console) as write_stdout: + write_stdout(f"current goal is {current_executing_goal.name!r}") + return OutputtingGoal(0) + + rule_runner = RuleRunner( + rules=collect_rules(locals()), + ) + result = rule_runner.run_goal_rule(OutputtingGoal) + assert result.exit_code == 0 + assert result.stdout == "current goal is 'dummy'" diff --git a/src/python/pants/engine/internals/scheduler.py b/src/python/pants/engine/internals/scheduler.py index dcaec6e6279..dc4c31265d1 100644 --- a/src/python/pants/engine/internals/scheduler.py +++ b/src/python/pants/engine/internals/scheduler.py @@ -32,7 +32,7 @@ Snapshot, SymlinkEntry, ) -from pants.engine.goal import Goal +from pants.engine.goal import CurrentExecutingGoal, Goal from pants.engine.internals import native_engine from pants.engine.internals.docker import DockerResolveImageRequest, DockerResolveImageResult from pants.engine.internals.native_engine import ( @@ -349,6 +349,9 @@ class SchedulerSession: def __init__(self, scheduler: Scheduler, session: PySession) -> None: self._scheduler = scheduler self._py_session = session + self._current_goal = ( + session.session_values.get(CurrentExecutingGoal) or CurrentExecutingGoal() + ) @property def scheduler(self) -> Scheduler: @@ -526,16 +529,19 @@ def run_goal_rule( :param poll_delay: See self.execution_request. :returns: An exit_code for the given Goal. """ - if self._scheduler.visualize_to_dir is not None: - rule_graph_name = f"rule_graph.{product.name}.dot" - params = self._scheduler._to_params_list(subject) - self._scheduler.visualize_rule_subgraph_to_file( - os.path.join(self._scheduler.visualize_to_dir, rule_graph_name), - [type(p) for p in params], - product, + with self._current_goal._executing(product): + if self._scheduler.visualize_to_dir is not None: + rule_graph_name = f"rule_graph.{product.name}.dot" + params = self._scheduler._to_params_list(subject) + self._scheduler.visualize_rule_subgraph_to_file( + os.path.join(self._scheduler.visualize_to_dir, rule_graph_name), + [type(p) for p in params], + product, + ) + (return_value,) = self.product_request( + product, [subject], poll=poll, poll_delay=poll_delay ) - (return_value,) = self.product_request(product, [subject], poll=poll, poll_delay=poll_delay) - return cast(int, return_value.exit_code) + return cast(int, return_value.exit_code) def product_request( self, diff --git a/src/python/pants/init/engine_initializer.py b/src/python/pants/init/engine_initializer.py index 01d498352dd..6d2c5927104 100644 --- a/src/python/pants/init/engine_initializer.py +++ b/src/python/pants/init/engine_initializer.py @@ -20,7 +20,7 @@ from pants.engine.console import Console from pants.engine.environment import EnvironmentName from pants.engine.fs import PathGlobs, Snapshot, Workspace -from pants.engine.goal import Goal +from pants.engine.goal import CurrentExecutingGoal, Goal from pants.engine.internals import ( build_files, dep_rules, @@ -263,6 +263,10 @@ def union_membership_singleton() -> UnionMembership: def build_root_singleton() -> BuildRoot: return cast(BuildRoot, BuildRoot.instance) + @rule + def current_executing_goal(session_values: SessionValues) -> CurrentExecutingGoal: + return session_values.get(CurrentExecutingGoal) or CurrentExecutingGoal() + # Create a Scheduler containing graph and filesystem rules, with no installed goals. rules = FrozenOrderedSet( ( diff --git a/src/python/pants/testutil/rule_runner.py b/src/python/pants/testutil/rule_runner.py index a5a027d3a2a..7a587b33c34 100644 --- a/src/python/pants/testutil/rule_runner.py +++ b/src/python/pants/testutil/rule_runner.py @@ -37,7 +37,7 @@ from pants.engine.env_vars import CompleteEnvironmentVars from pants.engine.environment import EnvironmentName from pants.engine.fs import Digest, PathGlobs, PathGlobsAndRoot, Snapshot, Workspace -from pants.engine.goal import Goal +from pants.engine.goal import CurrentExecutingGoal, Goal from pants.engine.internals import native_engine from pants.engine.internals.native_engine import ProcessConfigFromEnvironment, PyExecutor from pants.engine.internals.scheduler import ExecutionError, Scheduler, SchedulerSession @@ -314,6 +314,7 @@ def _set_new_session(self, scheduler: Scheduler) -> None: { OptionsBootstrapper: self.options_bootstrapper, CompleteEnvironmentVars: self.environment, + CurrentExecutingGoal: CurrentExecutingGoal(), **self.extra_session_values, } ), From de8391bb088b7bb540d1a64e770a8a17348dc007 Mon Sep 17 00:00:00 2001 From: Andreas Stenius Date: Tue, 11 Apr 2023 12:19:39 -0400 Subject: [PATCH 2/2] Move `CurrentExecutingGoal` to the goal param types. --- src/python/pants/bin/local_pants_runner.py | 2 -- src/python/pants/engine/goal.py | 16 ++---------- .../pants/engine/internals/scheduler.py | 26 +++++++------------ src/python/pants/init/engine_initializer.py | 16 +++++++----- src/python/pants/testutil/rule_runner.py | 2 +- 5 files changed, 23 insertions(+), 39 deletions(-) diff --git a/src/python/pants/bin/local_pants_runner.py b/src/python/pants/bin/local_pants_runner.py index 6874427191d..f732a01c071 100644 --- a/src/python/pants/bin/local_pants_runner.py +++ b/src/python/pants/bin/local_pants_runner.py @@ -13,7 +13,6 @@ from pants.build_graph.build_configuration import BuildConfiguration from pants.core.util_rules.environments import determine_bootstrap_environment from pants.engine.env_vars import CompleteEnvironmentVars -from pants.engine.goal import CurrentExecutingGoal from pants.engine.internals import native_engine from pants.engine.internals.native_engine import PyExecutor, PySessionCancellationLatch from pants.engine.internals.scheduler import ExecutionError @@ -144,7 +143,6 @@ def create( { OptionsBootstrapper: options_bootstrapper, CompleteEnvironmentVars: env, - CurrentExecutingGoal: CurrentExecutingGoal(), } ), cancellation_latch=cancellation_latch, diff --git a/src/python/pants/engine/goal.py b/src/python/pants/engine/goal.py index b1ee5a2df0f..bd3a508d02c 100644 --- a/src/python/pants/engine/goal.py +++ b/src/python/pants/engine/goal.py @@ -137,20 +137,8 @@ def name(cls) -> str: @dataclass(frozen=True) class CurrentExecutingGoal(EngineAwareReturnType): - goal: type[Goal] | None = None - - @property - def name(self) -> str | None: - return None if self.goal is None else self.goal.name - - @contextmanager - def _executing(self, goal: type[Goal]) -> Iterator[None]: - # Mutate current goal; we're only frozen to avoid inadvertent tampering with `self.goal`. - object.__setattr__(self, "goal", goal) - try: - yield - finally: - object.__setattr__(self, "goal", None) + name: str + goal: type[Goal] def cacheable(self) -> bool: return False diff --git a/src/python/pants/engine/internals/scheduler.py b/src/python/pants/engine/internals/scheduler.py index 18f5f878026..7d0d8117dd3 100644 --- a/src/python/pants/engine/internals/scheduler.py +++ b/src/python/pants/engine/internals/scheduler.py @@ -32,7 +32,7 @@ Snapshot, SymlinkEntry, ) -from pants.engine.goal import CurrentExecutingGoal, Goal +from pants.engine.goal import Goal from pants.engine.internals import native_engine from pants.engine.internals.docker import DockerResolveImageRequest, DockerResolveImageResult from pants.engine.internals.native_engine import ( @@ -354,9 +354,6 @@ class SchedulerSession: def __init__(self, scheduler: Scheduler, session: PySession) -> None: self._scheduler = scheduler self._py_session = session - self._current_goal = ( - session.session_values.get(CurrentExecutingGoal) or CurrentExecutingGoal() - ) @property def scheduler(self) -> Scheduler: @@ -534,19 +531,16 @@ def run_goal_rule( :param poll_delay: See self.execution_request. :returns: An exit_code for the given Goal. """ - with self._current_goal._executing(product): - if self._scheduler.visualize_to_dir is not None: - rule_graph_name = f"rule_graph.{product.name}.dot" - params = self._scheduler._to_params_list(subject) - self._scheduler.visualize_rule_subgraph_to_file( - os.path.join(self._scheduler.visualize_to_dir, rule_graph_name), - [type(p) for p in params], - product, - ) - (return_value,) = self.product_request( - product, [subject], poll=poll, poll_delay=poll_delay + if self._scheduler.visualize_to_dir is not None: + rule_graph_name = f"rule_graph.{product.name}.dot" + params = self._scheduler._to_params_list(subject) + self._scheduler.visualize_rule_subgraph_to_file( + os.path.join(self._scheduler.visualize_to_dir, rule_graph_name), + [type(p) for p in params], + product, ) - return cast(int, return_value.exit_code) + (return_value,) = self.product_request(product, [subject], poll=poll, poll_delay=poll_delay) + return cast(int, return_value.exit_code) def product_request( self, diff --git a/src/python/pants/init/engine_initializer.py b/src/python/pants/init/engine_initializer.py index 445d28d5636..631bae5b9bb 100644 --- a/src/python/pants/init/engine_initializer.py +++ b/src/python/pants/init/engine_initializer.py @@ -98,7 +98,13 @@ class GraphSession: goal_map: Any # NB: Keep this in sync with the method `run_goal_rules`. - goal_param_types: ClassVar[tuple[type, ...]] = (Specs, Console, Workspace, EnvironmentName) + goal_param_types: ClassVar[tuple[type, ...]] = ( + Specs, + Console, + Workspace, + EnvironmentName, + CurrentExecutingGoal, + ) def goal_consumed_subsystem_scopes(self, goal_name: str) -> tuple[str, ...]: """Return the scopes of subsystems that could be consumed while running the given goal.""" @@ -143,7 +149,9 @@ def run_goal_rules( if not goal_product.subsystem_cls.activated(union_membership): raise GoalNotActivatedException(goal) # NB: Keep this in sync with the property `goal_param_types`. - params = Params(specs, self.console, workspace, env_name) + params = Params( + specs, self.console, workspace, env_name, CurrentExecutingGoal(goal, goal_product) + ) logger.debug(f"requesting {goal_product} to satisfy execution of `{goal}` goal") try: exit_code = self.scheduler_session.run_goal_rule( @@ -268,10 +276,6 @@ def union_membership_singleton() -> UnionMembership: def build_root_singleton() -> BuildRoot: return cast(BuildRoot, BuildRoot.instance) - @rule - def current_executing_goal(session_values: SessionValues) -> CurrentExecutingGoal: - return session_values.get(CurrentExecutingGoal) or CurrentExecutingGoal() - # Create a Scheduler containing graph and filesystem rules, with no installed goals. rules = FrozenOrderedSet( ( diff --git a/src/python/pants/testutil/rule_runner.py b/src/python/pants/testutil/rule_runner.py index 66f07703469..04b568bed91 100644 --- a/src/python/pants/testutil/rule_runner.py +++ b/src/python/pants/testutil/rule_runner.py @@ -374,7 +374,6 @@ def _set_new_session(self, scheduler: Scheduler) -> None: { OptionsBootstrapper: self.options_bootstrapper, CompleteEnvironmentVars: self.environment, - CurrentExecutingGoal: CurrentExecutingGoal(), **self.extra_session_values, } ), @@ -437,6 +436,7 @@ def run_goal_rule( console, Workspace(self.scheduler), *([self.inherent_environment] if self.inherent_environment else []), + CurrentExecutingGoal(goal.name, goal), ), )