bxl: allow users to make atarget sets

app/buck2_build_api/src/query/bxl.rs

@@ -12,6 +12,7 @@ use std::future::Future;
 use std::pin::Pin;
 
 use async_trait::async_trait;
+use buck2_artifact::actions::key::ActionKey;
 use buck2_core::cells::CellResolver;
 use buck2_core::cells::name::CellName;
 use buck2_core::configuration::compatibility::MaybeCompatible;
@@ -183,6 +184,11 @@ pub trait BxlAqueryFunctions: Send {
         dice: &mut DiceComputations<'_>,
         targets: &TargetSet<ActionQueryNode>,
     ) -> buck2_error::Result<TargetSet<ActionQueryNode>>;
+    async fn get_action_nodes(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        action_keys: Vec<ActionKey>,
+    ) -> buck2_error::Result<TargetSet<ActionQueryNode>>;
 }
 
 pub static NEW_BXL_CQUERY_FUNCTIONS: LateBinding<

app/buck2_bxl/src/bxl/starlark_defs/analysis_result.rs

@@ -15,6 +15,7 @@ use buck2_build_api::analysis::AnalysisResult;
 use buck2_build_api::interpreter::rule_defs::provider::collection::FrozenProviderCollection;
 use buck2_build_api::interpreter::rule_defs::provider::dependency::Dependency;
 use buck2_core::provider::label::ConfiguredProvidersLabel;
+use buck2_interpreter::types::configured_providers_label::StarlarkConfiguredProvidersLabel;
 use dupe::Dupe;
 use starlark::any::ProvidesStaticType;
 use starlark::environment::Methods;
@@ -100,6 +101,22 @@ fn starlark_analysis_result_methods(builder: &mut MethodsBuilder) {
         }
     }
 
+    /// Gets the configured providers label for this analysis result.
+    ///
+    /// Sample usage:
+    /// ```python
+    /// def _impl_label(ctx):
+    ///     actions = ctx.aquery().all_actions("//target")
+    ///     for node in actions:
+    ///         if analysis := node.analysis():
+    ///             ctx.output.print(analysis.label())
+    /// ```
+    fn label<'v>(
+        this: &StarlarkAnalysisResult,
+    ) -> starlark::Result<StarlarkConfiguredProvidersLabel> {
+        Ok(StarlarkConfiguredProvidersLabel::new(this.label.dupe()))
+    }
+
     /// Converts the analysis result into a `Dependency`. Currently, you can only get a `Dependency` without any
     /// transitions. This means that you cannot create an exec dep or toolchain from an analysis result.
     ///

app/buck2_bxl/src/bxl/starlark_defs/aquery.rs

@@ -9,6 +9,7 @@
  */
 
 use allocative::Allocative;
+use buck2_artifact::actions::key::ActionKey;
 use buck2_build_api::actions::query::ActionQueryNode;
 use buck2_build_api::query::bxl::BxlAqueryFunctions;
 use buck2_build_api::query::bxl::NEW_BXL_AQUERY_FUNCTIONS;
@@ -46,6 +47,7 @@ use starlark::values::type_repr::StarlarkTypeRepr;
 
 use crate::bxl::starlark_defs::context::BxlContext;
 use crate::bxl::starlark_defs::context::ErrorPrinter;
+use crate::bxl::starlark_defs::nodes::action::StarlarkAction;
 use crate::bxl::starlark_defs::nodes::action::StarlarkActionQueryNode;
 use crate::bxl::starlark_defs::providers_expr::AnyProvidersExprArg;
 use crate::bxl::starlark_defs::providers_expr::ProvidersExpr;
@@ -122,6 +124,7 @@ enum UnpackActionNodes<'v> {
     ActionQueryNodesSet(&'v StarlarkTargetSet<ActionQueryNode>),
     ConfiguredProviders(AnyProvidersExprArg<'v>),
     ConfiguredTargets(ConfiguredTargetListExprArg<'v>),
+    StarlarkActions(UnpackList<StarlarkAction>),
 }
 
 // Aquery operates on `ActionQueryNode`s. Under the hood, the target set of action query nodes is obtained
@@ -140,6 +143,14 @@ async fn unpack_action_nodes<'v>(
             return Ok(action_nodes.into_iter().map(|v| v.0).collect());
         }
         UnpackActionNodes::ActionQueryNodesSet(action_nodes) => return Ok(action_nodes.0.clone()),
+        UnpackActionNodes::StarlarkActions(actions) => {
+            let action_keys: Vec<ActionKey> = actions
+                .into_iter()
+                .map(|starlark_action| starlark_action.0.key().dupe())
+                .collect();
+
+            return aquery_env.get_action_nodes(dice, action_keys).await;
+        }
         UnpackActionNodes::ConfiguredProviders(arg) => {
             ProvidersExpr::<ConfiguredProvidersLabel>::unpack(
                 arg,
@@ -183,6 +194,11 @@ async fn unpack_action_nodes<'v>(
 ///
 /// Query results are `target_set`s of `action_query_node`s, which supports iteration,
 /// indexing, `len()`, set addition/subtraction, and `equals()`.
+///
+/// Actions can be specified as:
+/// - Target expressions (configured targets/providers)
+/// - Existing action query nodes or target sets
+/// - `bxl.Action` objects (obtained from `ctx.audit().output()`)
 #[starlark_module]
 fn aquery_methods(builder: &mut MethodsBuilder) {
     /// The deps query for finding the transitive closure of dependencies.

app/buck2_bxl/src/bxl/starlark_defs/functions.rs

@@ -10,6 +10,7 @@
 
 use std::time::Instant;
 
+use buck2_build_api::actions::query::ActionQueryNode;
 use buck2_build_api::interpreter::rule_defs::artifact::starlark_artifact_like::ValueAsInputArtifactLikeUnpack;
 use buck2_build_api::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
 use buck2_core::cells::CellAliasResolver;
@@ -42,6 +43,7 @@ use super::context::output::get_cmd_line_inputs;
 use super::nodes::unconfigured::StarlarkTargetNode;
 use crate::bxl::starlark_defs::context::BxlContext;
 use crate::bxl::starlark_defs::eval_extra::BxlEvalExtra;
+use crate::bxl::starlark_defs::nodes::action::StarlarkActionQueryNode;
 use crate::bxl::starlark_defs::nodes::configured::StarlarkConfiguredTargetNode;
 use crate::bxl::starlark_defs::targetset::StarlarkTargetSet;
 use crate::bxl::starlark_defs::time::StarlarkInstant;
@@ -92,6 +94,30 @@ pub(crate) fn register_target_function(builder: &mut GlobalsBuilder) {
                 .map(|node| node.0),
         ))
     }
+
+    /// Creates a target set from a list of action query nodes.
+    ///
+    /// Sample usage:
+    /// ```python
+    /// def _impl_atarget_set(ctx):
+    ///     actions = ctx.aquery().all_actions("//target")
+    ///     action_a = actions[0]
+    ///     action_b = actions[1]
+    ///     action_set = bxl.atarget_set([action_a, action_b])
+    ///     # Now can use in further queries
+    ///     deps = ctx.aquery().deps(action_set)
+    /// ```
+    fn atarget_set(
+        nodes: Option<UnpackList<StarlarkActionQueryNode>>,
+    ) -> starlark::Result<StarlarkTargetSet<ActionQueryNode>> {
+        Ok(StarlarkTargetSet::from_iter(
+            nodes
+                .unwrap_or(UnpackList::default())
+                .items
+                .into_iter()
+                .map(|node| node.0),
+        ))
+    }
 }
 
 #[derive(Debug, buck2_error::Error, Clone)]

app/buck2_bxl/src/bxl/starlark_defs/nodes/action.rs

@@ -12,9 +12,11 @@ use std::convert::Infallible;
 use std::sync::Arc;
 
 use allocative::Allocative;
+use buck2_artifact::artifact::artifact_type::Artifact;
 use buck2_build_api::actions::RegisteredAction;
 use buck2_build_api::actions::query::ActionQueryNode;
 use buck2_build_api::actions::query::OwnedActionAttr;
+use buck2_build_api::interpreter::rule_defs::artifact::starlark_artifact::StarlarkArtifact;
 use buck2_core::deferred::base_deferred_key::BaseDeferredKey;
 use buck2_error::buck2_error;
 use buck2_interpreter::types::target_label::StarlarkConfiguredTargetLabel;
@@ -49,7 +51,7 @@ pub(crate) struct StarlarkAction(pub(crate) Arc<RegisteredAction>);
 
 starlark_simple_value!(StarlarkAction);
 
-#[starlark_value(type = "action")]
+#[starlark_value(type = "bxl.Action")]
 impl<'v> StarlarkValue<'v> for StarlarkAction {
     fn get_methods() -> Option<&'static Methods> {
         static RES: MethodsStatic = MethodsStatic::new();
@@ -67,7 +69,7 @@ impl<'a> UnpackValue<'a> for StarlarkAction {
     }
 }
 
-/// Methods for an action.
+/// Methods for an action obtained from [`bxl.AuditContext.output()`](../AuditContext#output).
 #[starlark_module]
 fn action_methods(builder: &mut MethodsBuilder) {
     /// Gets the owning configured target label for an action.
@@ -90,6 +92,17 @@ fn action_methods(builder: &mut MethodsBuilder) {
             .into()),
         }
     }
+
+    /// Gets the artifacts built by this action.
+    fn outputs<'v>(this: StarlarkAction) -> starlark::Result<Vec<StarlarkArtifact>> {
+        Ok(this
+            .0
+            .action()
+            .outputs()
+            .iter()
+            .map(|a| StarlarkArtifact::new(Artifact::from(a.dupe())))
+            .collect())
+    }
 }
 
 #[derive(Debug, Display, ProvidesStaticType, Allocative)]

app/buck2_query_impls/src/aquery/bxl.rs

@@ -12,6 +12,7 @@ use std::marker::PhantomData;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use buck2_artifact::actions::key::ActionKey;
 use buck2_build_api::actions::query::ActionQueryNode;
 use buck2_build_api::analysis::calculation::RuleAnalysisCalculation;
 use buck2_build_api::query::bxl::BxlAqueryFunctions;
@@ -303,6 +304,24 @@ impl BxlAqueryFunctions for BxlAqueryFunctionsImpl {
         })
         .await
     }
+
+    async fn get_action_nodes(
+        &self,
+        dice: &mut DiceComputations<'_>,
+        action_keys: Vec<ActionKey>,
+    ) -> buck2_error::Result<TargetSet<ActionQueryNode>> {
+        dice.with_linear_recompute(|dice| async move {
+            let delegate = self.aquery_delegate(&dice).await?;
+            let mut result = TargetSet::new();
+            let nodes = buck2_util::future::try_join_all(
+                action_keys.iter().map(|key| delegate.get_action_node(&key)),
+            )
+            .await?;
+            result.extend(nodes);
+            Ok(result)
+        })
+        .await
+    }
 }
 
 pub(crate) fn init_new_bxl_aquery_functions() {

tests/core/bxl/test_audit_data/audit.bxl

@@ -15,6 +15,7 @@ def _audit_output_action_exists_impl(ctx):
     action = ctx.audit().output(buck_out)
 
     asserts.equals(action.owner(), target.label)
+    asserts.equals(action.outputs()[0].basename, "with_output.txt")
 
 audit_output_action_exists = bxl_main(
     impl = _audit_output_action_exists_impl,

tests/core/query/aquery/test_aquery.py

@@ -154,3 +154,28 @@ async def test_bxl_aquery_eval(buck: Buck) -> None:
 @buck_test()
 async def test_bxl_aquery_action_query_node(buck: Buck) -> None:
     await buck.bxl("//:aquery.bxl:action_query_node")
+
+
+# Tests for bxl.Action support in aquery operations
+@buck_test()
+async def test_bxl_action_deps_0(buck: Buck) -> None:
+    """Test passing a bxl.Action to aquery.deps() with depth=0 returns itself"""
+    await buck.bxl("//:aquery.bxl:action_deps_0")
+
+
+@buck_test()
+async def test_bxl_action_deps(buck: Buck) -> None:
+    """Test that you can isolate the deps of one action by itself"""
+    await buck.bxl("//:aquery.bxl:action_deps")
+
+
+@buck_test()
+async def test_bxl_atarget_set(buck: Buck) -> None:
+    """Test bxl.atarget_set() creates a target set from ActionQueryNodes"""
+    await buck.bxl("//:aquery.bxl:atarget_set_test")
+
+
+@buck_test()
+async def test_bxl_analysis_label(buck: Buck) -> None:
+    """Test analysis.label() returns the configured providers label"""
+    await buck.bxl("//:aquery.bxl:analysis_label_test")

tests/core/query/aquery/test_aquery_data/aquery.bxl

@@ -130,7 +130,7 @@ def _impl_action_query_node(ctx):
     action = result[0]
     analysis = result[1]
 
-    _assert_eq(type(action.action()), "action")
+    _assert_eq(type(action.action()), "bxl.Action")
     _assert_eq(action.rule_type, "copy")
     _assert_eq(str(action.action().owner().raw_target()), "root//:test")
 
@@ -141,3 +141,91 @@ action_query_node = bxl_main(
     impl = _impl_action_query_node,
     cli_args = {},
 )
+
+def _impl_action_deps_0(ctx):
+    """Test passing a bxl.Action to aquery.deps() returns itself"""
+    nodes = ctx.aquery().eval("//:test")
+
+    action = nodes[0].action()
+
+    # Pass the bxl.Action rather than a target to aquery.deps()
+    result = ctx.aquery().deps([action], depth = 0)
+    _assert_eq(action.outputs(), result[0].action().outputs())
+
+action_deps_0 = bxl_main(
+    impl = _impl_action_deps_0,
+    cli_args = {},
+)
+
+def _impl_action_deps(ctx):
+    """Test that you can isolate the deps of one action by itself"""
+    nodes = ctx.aquery().eval("//:test")
+
+    action = nodes[0].action()
+
+    # Pass the bxl.Action rather than a target to aquery.deps()
+    result = ctx.aquery().deps([action], depth = 1)
+
+    _assert_eq(len(result), 2)
+    # First should be itself
+    _assert_eq(action.outputs(), result[0].action().outputs())
+    dep = result[1]
+    _assert_eq(dep.action().outputs()[0].short_path, "dep")
+
+action_deps = bxl_main(
+    impl = _impl_action_deps,
+    cli_args = {},
+)
+
+def _impl_atarget_set_test(ctx):
+    """Test bxl.atarget_set() creates a target set from ActionQueryNodes"""
+    nodes = ctx.aquery().eval("//:test")
+
+    # Get two action query nodes
+    node_a = nodes[0]
+    node_b = nodes[1]
+
+    # Create a target set from action query nodes
+    action_set = bxl.atarget_set([node_a, node_b])
+
+    # Verify it's a target_set type
+    _assert_eq(type(action_set), "target_set")
+
+    # Verify we can use it in further aquery operations
+    deps = ctx.aquery().deps(action_set, depth = 0)
+    _assert_eq(len(deps), 2)
+
+atarget_set_test = bxl_main(
+    impl = _impl_atarget_set_test,
+    cli_args = {},
+)
+
+def _impl_analysis_label_test(ctx):
+    """Test analysis.label() returns the configured providers label"""
+    nodes = ctx.aquery().eval("//:test")
+
+    # Find an analysis node by iterating through results
+    analysis = None
+    for node in nodes:
+        if node.rule_type == "analysis":
+            analysis = node.analysis()
+            break
+
+    if not analysis:
+        fail("No analysis node found in results")
+
+    _assert_eq(type(analysis), "bxl.AnalysisResult")
+
+    # Test the label() method
+    label = analysis.label()
+    _assert_eq(type(label), "Label")
+
+    # Verify it's a valid label (should contain the target)
+    label_str = str(label)
+    if "root//:test" not in label_str:
+        fail("Expected label to contain 'root//:test', got: {}".format(label_str))
+
+analysis_label_test = bxl_main(
+    impl = _impl_analysis_label_test,
+    cli_args = {},
+)

tests/core/query/aquery/test_aquery_data/prelude.bzl

@@ -9,6 +9,7 @@
 def _test(ctx: AnalysisContext):
     dep = ctx.actions.write("dep", "")
     default = ctx.actions.copy_file("default", dep)
+    other_default = ctx.actions.copy_file("other_default", dep)
     other = ctx.actions.write("other", "")
 
     sub_default = ctx.actions.write("sub_default", "")
@@ -18,7 +19,7 @@ def _test(ctx: AnalysisContext):
     ctx.actions.write("unused", "")
 
     return [DefaultInfo(
-        default_outputs = [default],
+        default_outputs = [default, other_default],
         other_outputs = [other],
         sub_targets = {
             "sub": [