add union to gather BSP BuildTarget instances

[ci skip-rust]

[ci skip-build-wheels]

src/python/pants/bsp/protocol.py

@@ -15,7 +15,12 @@
 )
 from pylsp_jsonrpc.streams import JsonRpcStreamReader, JsonRpcStreamWriter  # type: ignore[import]
 
-from pants.bsp.spec import InitializeBuildParams, InitializeBuildResult
+from pants.bsp.spec import (
+    InitializeBuildParams,
+    InitializeBuildResult,
+    WorkspaceBuildTargetsParams,
+    WorkspaceBuildTargetsResult,
+)
 from pants.engine.internals.scheduler import SchedulerSession
 from pants.util.frozendict import FrozenDict
 
@@ -33,6 +38,9 @@ class _HandlerMapping:
     "build/initialize": _HandlerMapping(
         InitializeBuildParams.from_json_dict, InitializeBuildResult
     ),
+    "workspace/buildTargets": _HandlerMapping(
+        WorkspaceBuildTargetsParams.from_json_dict, WorkspaceBuildTargetsResult
+    ),
 }
 
 
@@ -97,6 +105,8 @@ def _handle_inbound_message(self, *, method_name: str, params: Any):
         )
         returns, throws = self._scheduler_session.execute(execution_request)
         if len(returns) == 1 and len(throws) == 0:
+            if method_name == self._INITIALIZE_METHOD_NAME:
+                self._initialized = True
             return returns[0][1].value.to_json_dict()
         elif len(returns) == 0 and len(throws) == 1:
             raise throws[0][1].exc

src/python/pants/bsp/protocol_test.py

@@ -3,7 +3,6 @@
 import os
 from contextlib import contextmanager
 from dataclasses import dataclass
-from pathlib import Path
 from threading import Thread
 from typing import BinaryIO
 
@@ -14,8 +13,14 @@
 
 from pants.bsp.protocol import BSPConnection
 from pants.bsp.rules import rules as bsp_rules
-from pants.bsp.spec import BuildClientCapabilities, InitializeBuildParams, InitializeBuildResult
-from pants.engine.internals.scheduler_test_base import SchedulerTestBase
+from pants.bsp.spec import (
+    BuildClientCapabilities,
+    InitializeBuildParams,
+    InitializeBuildResult,
+    WorkspaceBuildTargetsParams,
+    WorkspaceBuildTargetsResult,
+)
+from pants.testutil.rule_runner import RuleRunner
 
 
 @dataclass(frozen=True)
@@ -52,48 +57,55 @@ def setup_pipes():
         outbound_writer.close()
 
 
-class TestBSPConnection(SchedulerTestBase):
-    def test_errors_for_uninitialized_connection(self, tmp_path: Path) -> None:
-        with setup_pipes() as pipes:
-            # TODO: This code should be moved to a context manager. For now, only the pipes are managed
-            # with a context manager.
-            scheduler = self.mk_scheduler(tmp_path, [*bsp_rules()])
-            conn = BSPConnection(scheduler, pipes.inbound_reader, pipes.outbound_writer)
-
-            def run_bsp_server():
-                conn.run()
-
-            bsp_thread = Thread(target=run_bsp_server)
-            bsp_thread.daemon = True
-            bsp_thread.start()
-
-            client_reader = JsonRpcStreamReader(pipes.outbound_reader)
-            client_writer = JsonRpcStreamWriter(pipes.inbound_writer)
-            endpoint = Endpoint({}, lambda msg: client_writer.write(msg))
-
-            def run_client():
-                client_reader.listen(lambda msg: endpoint.consume(msg))
-
-            client_thread = Thread(target=run_client)
-            client_thread.daemon = True
-            client_thread.start()
-
-            response_fut = endpoint.request("foo")
-            with pytest.raises(JsonRpcException) as exc_info:
-                response_fut.result(timeout=15)
-            assert exc_info.value.code == -32002
-            assert exc_info.value.message == "Client must first call `build/initialize`."
-
-            init_request = InitializeBuildParams(
-                display_name="test",
-                version="0.0.0",
-                bsp_version="0.0.0",
-                root_uri="https://example.com",
-                capabilities=BuildClientCapabilities(language_ids=()),
-                data=None,
-            )
-            response_fut = endpoint.request("build/initialize", init_request.to_json_dict())
-            raw_response = response_fut.result(timeout=15)
-            response = InitializeBuildResult.from_json_dict(raw_response)
-            assert response.display_name == "Pants"
-            assert response.bsp_version == "0.0.1"
+def test_basic_bsp_protocol() -> None:
+    with setup_pipes() as pipes:
+        # TODO: This code should be moved to a context manager. For now, only the pipes are managed
+        # with a context manager.
+        rule_runner = RuleRunner(rules=bsp_rules())
+        conn = BSPConnection(rule_runner.scheduler, pipes.inbound_reader, pipes.outbound_writer)
+
+        def run_bsp_server():
+            conn.run()
+
+        bsp_thread = Thread(target=run_bsp_server)
+        bsp_thread.daemon = True
+        bsp_thread.start()
+
+        client_reader = JsonRpcStreamReader(pipes.outbound_reader)
+        client_writer = JsonRpcStreamWriter(pipes.inbound_writer)
+        endpoint = Endpoint({}, lambda msg: client_writer.write(msg))
+
+        def run_client():
+            client_reader.listen(lambda msg: endpoint.consume(msg))
+
+        client_thread = Thread(target=run_client)
+        client_thread.daemon = True
+        client_thread.start()
+
+        response_fut = endpoint.request("foo")
+        with pytest.raises(JsonRpcException) as exc_info:
+            response_fut.result(timeout=15)
+        assert exc_info.value.code == -32002
+        assert exc_info.value.message == "Client must first call `build/initialize`."
+
+        init_request = InitializeBuildParams(
+            display_name="test",
+            version="0.0.0",
+            bsp_version="0.0.0",
+            root_uri="https://example.com",
+            capabilities=BuildClientCapabilities(language_ids=()),
+            data=None,
+        )
+        response_fut = endpoint.request("build/initialize", init_request.to_json_dict())
+        raw_response = response_fut.result(timeout=15)
+        response = InitializeBuildResult.from_json_dict(raw_response)
+        assert response.display_name == "Pants"
+        assert response.bsp_version == "0.0.1"
+
+        build_targets_request = WorkspaceBuildTargetsParams()
+        response_fut = endpoint.request(
+            "workspace/buildTargets", build_targets_request.to_json_dict()
+        )
+        raw_response = response_fut.result(timeout=15)
+        response = WorkspaceBuildTargetsResult.from_json_dict(raw_response)
+        assert response.targets == ()

src/python/pants/bsp/rules.py

@@ -1,10 +1,34 @@
 # Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
-from pants.bsp.spec import BuildServerCapabilities, InitializeBuildParams, InitializeBuildResult
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from pants.bsp.spec import (
+    BuildServerCapabilities,
+    BuildTarget,
+    InitializeBuildParams,
+    InitializeBuildResult,
+    WorkspaceBuildTargetsParams,
+    WorkspaceBuildTargetsResult,
+)
+from pants.engine.internals.selectors import Get, MultiGet
 from pants.engine.rules import QueryRule, collect_rules, rule
+from pants.engine.unions import UnionMembership, union
 from pants.version import VERSION
 
 
+@union
+class BSPBuildTargetsRequest:
+    """Request language backends to provide BSP `BuildTarget` instances for their managed target
+    types."""
+
+
+@dataclass(frozen=True)
+class BSPBuildTargets:
+    targets: tuple[BuildTarget, ...]
+
+
 @rule
 async def bsp_build_initialize(_request: InitializeBuildParams) -> InitializeBuildResult:
     return InitializeBuildResult(
@@ -27,8 +51,27 @@ async def bsp_build_initialize(_request: InitializeBuildParams) -> InitializeBui
     )
 
 
+@rule
+async def bsp_workspace_build_targets(
+    _: WorkspaceBuildTargetsParams, union_membership: UnionMembership
+) -> WorkspaceBuildTargetsResult:
+    request_types = union_membership.get(BSPBuildTargetsRequest)
+    responses = await MultiGet(
+        Get(BSPBuildTargets, BSPBuildTargetsRequest, request_type())
+        for request_type in request_types
+    )
+    result: list[BuildTarget] = []
+    for response in responses:
+        result.extend(response.targets)
+    result.sort(key=lambda btgt: btgt.id.uri)
+    return WorkspaceBuildTargetsResult(
+        targets=tuple(result),
+    )
+
+
 def rules():
     return (
         *collect_rules(),
         QueryRule(InitializeBuildResult, (InitializeBuildParams,)),
+        QueryRule(WorkspaceBuildTargetsResult, (WorkspaceBuildTargetsParams,)),
     )

src/python/pants/bsp/spec.py

@@ -28,6 +28,9 @@ class BuildTargetIdentifier:
     def from_json_dict(cls, d):
         return cls(uri=d["uri"])
 
+    def to_json_dict(self):
+        return {"uri": self.uri}
+
 
 @dataclass(frozen=True)
 class BuildTargetCapabilities:
@@ -52,6 +55,14 @@ def from_json_dict(cls, d):
             can_debug=d["canDebug"],
         )
 
+    def to_json_dict(self):
+        return {
+            "canCompile": self.can_compile,
+            "canTest": self.can_test,
+            "canRun": self.can_run,
+            "canDebug": self.can_debug,
+        }
+
 
 # Note: The BSP "build target" concept is _not_ the same as a Pants "target". They are similar but
 # should be not be conflated with one another.
@@ -98,6 +109,7 @@ class BuildTarget:
 
     # Language-specific metadata about this target.
     # See ScalaBuildTarget as an example.
+    # TODO: Figure out generic decode/encode of this field. Maybe use UnionRule to allow language backends to hook?
     data: Any | None
 
     @classmethod
@@ -116,6 +128,24 @@ def from_json_dict(cls, d):
             data=d.get("data"),
         )
 
+    def to_json_dict(self):
+        result = {
+            "id": self.id.to_json_dict(),
+            "capabilities": self.capabilities.to_json_dict(),
+            "tags": self.tags,
+            "languageIds": self.language_ids,
+            "dependencies": [dep.to_json_dict() for dep in self.dependencies],
+        }
+        if self.display_name is not None:
+            result["displayName"] = self.display_name
+        if self.base_directory is not None:
+            result["baseDirectory"] = self.base_directory
+        if self.data_kind is not None:
+            result["dataKind"] = self.data_kind
+        if self.data is not None:
+            result["data"] = self.data
+        return result
+
 
 class BuildTargetDataKind:
     # The `data` field contains a `ScalaBuildTarget` object.
@@ -442,3 +472,25 @@ def to_json_dict(self):
             # TODO: Figure out whether to encode/decode data in a generic manner.
             result["data"] = self.data
         return result
+
+
+@dataclass(frozen=True)
+class WorkspaceBuildTargetsParams:
+    @classmethod
+    def from_json_dict(cls, _d):
+        return cls()
+
+    def to_json_dict(self):
+        return {}
+
+
+@dataclass(frozen=True)
+class WorkspaceBuildTargetsResult:
+    targets: tuple[BuildTarget, ...]
+
+    @classmethod
+    def from_json_dict(cls, d):
+        return cls(targets=tuple(BuildTarget.from_json_dict(tgt) for tgt in d["targets"]))
+
+    def to_json_dict(self):
+        return {"targets": [tgt.to_json_dict() for tgt in self.targets]}