feat: Add utility to generate plan cache stubs

pyproject.toml

@@ -37,6 +37,7 @@ dependencies = [
     "graypy>=2.1.0",
     "httpx>=0.28.1",
     "aiohttp>=3.13.5",
+    "jinja2>=3.1.6",
 ]
 dynamic = ["version"]
 license.file = "LICENSE"

src/blueapi/cli/cli.py

@@ -239,6 +239,19 @@ def start_application(obj: dict):
     start(config)
 
 
+@main.command()
+@click.pass_obj
+@click.argument("target", type=click.Path(file_okay=False))
+def generate_stubs(obj: dict, target: Path):
+    click.echo(f"Writing stubs to {target}")
+
+    config: ApplicationConfig = obj["config"]
+    bc = BlueapiClient.from_config(config)
+    from . import stubgen
+
+    stubgen.generate_stubs(Path(target), list(bc.plans), list(bc.devices))
+
+
 @main.group()
 @click.option(
     "-o",

src/blueapi/cli/stubgen.py

@@ -0,0 +1,124 @@
+import logging
+from dataclasses import dataclass
+from inspect import cleandoc
+from pathlib import Path
+from textwrap import dedent
+from typing import Any, Self, TextIO
+
+from jinja2 import Environment, PackageLoader
+
+from blueapi.client.device_cache import DeviceRef
+from blueapi.client.plan_cache import Plan
+from blueapi.core import context
+from blueapi.core.bluesky_types import BLUESKY_PROTOCOLS
+
+log = logging.getLogger(__name__)
+
+PLAN_TMPL = "plan_cache_template.pyi.jinja"
+DEV_TMPL = "device_cache_template.pyi.jinja"
+
+
+@dataclass
+class ArgSpec:
+    name: str
+    type: str
+    optional: bool
+
+
+@dataclass
+class PlanSpec:
+    name: str
+    docs: str
+    args: list[ArgSpec]
+
+    @classmethod
+    def from_plan(cls, plan: Plan) -> Self:
+        req = set(plan.required)
+        args = [
+            ArgSpec(arg, _type_string(spec), arg not in req)
+            for arg, spec in plan.model.parameter_schema.get("properties", {}).items()
+        ]
+        return cls(plan.name, plan.help_text, args)
+
+
+BLUESKY_PROTOCOL_NAMES = {context.qualified_name(proto) for proto in BLUESKY_PROTOCOLS}
+
+
+def _type_string(spec) -> str:
+    """Best effort attempt at making useful type hints for plans"""
+    match spec.get("type"):
+        case "array" if "items" in spec:
+            return f"list[{_type_string(spec.get('items'))}]"
+        case "array" if "prefixItems" in spec:
+            items = spec.get("prefixItems")
+            return f"tuple[{', '.join(_type_string(item) for item in items)}]"
+        case "integer":
+            return "int"
+        case "number":
+            return "float"
+        case proto if proto in BLUESKY_PROTOCOL_NAMES:
+            return "DeviceRef"
+        case "object":
+            return "dict[str, Any]"
+        case "string":
+            return "str"
+        case "boolean":
+            return "bool"
+        case None if opts := spec.get("anyOf"):
+            return " | ".join(_type_string(opt) for opt in opts)
+        case _:
+            return "Any"
+
+
+def generate_stubs(target: Path, plans: list[Plan], devices: list[DeviceRef]):
+    log.info("Generating stubs for %d plans and %d devices", len(plans), len(devices))
+    target.mkdir(parents=True, exist_ok=True)
+    client_dir = target / "src" / "blueapi-stubs" / "client"
+
+    log.debug("Making project structure: %s", client_dir)
+    client_dir.mkdir(parents=True, exist_ok=True)
+
+    project_file = target / "pyproject.toml"
+    py_typed = target / "src" / "blueapi-stubs" / "py.typed"
+
+    log.debug("Writing pyproject.toml to %s", project_file)
+    with open(project_file, "w") as out:
+        out.write(
+            dedent("""
+                [project]
+                name = "blueapi-stubs"
+                version = "0.1.0"
+                description = "Generated client stubs for a running server"
+                readme = "README.md"
+                requires-python = ">=3.11"
+
+                dependencies = [
+                    "blueapi"
+                ]
+                """)
+        )
+
+    log.debug("Writing py.typed file to %s", py_typed)
+    with open(py_typed, "w") as out:
+        out.write("partial\n")
+
+    log.debug("Writing stub files to %s", client_dir)
+
+    plan_model = [PlanSpec.from_plan(p) for p in plans]
+    with open(client_dir / "plan_cache.pyi", "w") as out:
+        render_stub_file(out, PLAN_TMPL, plans=plan_model)
+
+    with open(client_dir / "device_cache.pyi", "w") as out:
+        render_stub_file(out, DEV_TMPL, devices=devices)
+
+
+def _docstring(text: str) -> str:
+    # """Convert a docstring to a format that can be inserted into the template"""
+    return cleandoc(text).replace('"""', '\\"""')
+
+
+def render_stub_file(stub_file: TextIO, template: str, **model: Any):
+    env = Environment(loader=PackageLoader("blueapi", package_path="stubs/templates"))
+    env.filters["docstring"] = _docstring
+    tmpl = env.get_template(template)
+    stub_file.write(tmpl.render(**model))

src/blueapi/client/client.py

@@ -1,12 +1,12 @@
 import itertools
 import logging
 import time
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from concurrent.futures import Future
 from contextlib import suppress
 from functools import cached_property
 from pathlib import Path
-from typing import Self
+from typing import Any, Self
 
 from bluesky_stomp.messaging import MessageContext, StompClient
 from bluesky_stomp.models import Broker
@@ -43,7 +43,6 @@
 from .device_cache import DeviceCache
 from .event_bus import AnyEvent, EventBusClient, OnAnyEvent
 from .plan_cache import PlanCache
-from .protocols import ClientProtocol
 from .rest import (
     BlueapiRestClient,
     BlueskyRemoteControlError,
@@ -61,7 +60,7 @@ class MissingInstrumentSessionError(Exception):
     pass
 
 
-class BlueapiClient(ClientProtocol):
+class BlueapiClient:
     """Unified client for controlling blueapi"""
 
     _rest: BlueapiRestClient
@@ -120,7 +119,7 @@ def for_host(cls, host: str | None = None) -> Self:
     @cached_property
     @start_as_current_span(TRACER)
     def plans(self) -> PlanCache:
-        return PlanCache(self, self._rest.get_plans().plans)
+        return PlanCache(self.run_plan, self._rest.get_plans().plans)
 
     @cached_property
     @start_as_current_span(TRACER)
@@ -321,6 +320,18 @@ def get_active_task(self) -> WorkerTask:
 
         return self.active_task
 
+    @start_as_current_span(TRACER, "name", "args", "kwargs")
+    def run_plan(
+        self, name: str, args: tuple[Any, ...], kwargs: Mapping[str, Any]
+    ) -> TaskStatus:
+        req = TaskRequest(
+            name=name,
+            args=list(args),
+            params=kwargs,
+            instrument_session=self.instrument_session,
+        )
+        return self.run_task(req)
+
     @start_as_current_span(TRACER, "task")
     def run_task(
         self,

src/blueapi/client/device_cache.py

@@ -1,13 +1,15 @@
+from blueapi.client.rest import BlueapiRestClient, ClientObjectRef
 from blueapi.service.model import DeviceModel
 
-from .protocols import ClientObjectRef, RestClientProtocol
 from .rest import NotFoundError
 
+# This file should be kept in sync with the type stub template in stubs/templates
+
 
 class DeviceCache:
     """Cache and lazily resolve devices from a Blueapi server."""
 
-    def __init__(self, rest: RestClientProtocol):
+    def __init__(self, rest: BlueapiRestClient):
         self._rest = rest
         self._cache = {
             model.name: DeviceRef(name=model.name, cache=self, model=model)

src/blueapi/client/plan_cache.py

@@ -1,12 +1,11 @@
 import logging
+from collections.abc import Callable
 from typing import Any
 
 from blueapi.client.rest import InvalidParametersError
-from blueapi.service.model import PlanModel, TaskRequest
+from blueapi.service.model import PlanModel
 from blueapi.worker.event import TaskError, TaskResult, TaskStatus
 
-from .protocols import ClientProtocol
-
 _REPR_MAX_LENGTH = 100
 _REPR_MAX_ARGS_INLINE = 3
 _JSON_TYPE_MAP = {
@@ -19,6 +18,10 @@
 
 log = logging.getLogger(__name__)
 
+# This file should be kept in sync with the type stub template in stubs/templates
+
+PlanRunner = Callable[[str, tuple[Any, ...], dict[str, Any]], TaskStatus]
+
 
 class PlanFailedError(Exception):
     def __init__(self, typ: str, message: str):
@@ -29,9 +32,9 @@ def __init__(self, typ: str, message: str):
 class PlanCache:
     """Collection of plans that can be accessed by name or attribute."""
 
-    def __init__(self, client: ClientProtocol, plans: list[PlanModel]):
+    def __init__(self, runner: PlanRunner, plans: list[PlanModel]):
         self._cache = {
-            model.name: Plan(name=model.name, model=model, client=client)
+            model.name: Plan(name=model.name, model=model, runner=runner)
             for model in plans
         }
         for name, plan in self._cache.items():
@@ -55,10 +58,10 @@ def __repr__(self) -> str:
 class Plan:
     """Callable client-side reference to a registered BlueAPI plan."""
 
-    def __init__(self, name: str, model: PlanModel, client: ClientProtocol):
+    def __init__(self, name: str, model: PlanModel, runner: PlanRunner):
         self.name = name
         self.model = model
-        self._client = client
+        self._runner = runner
         self.__doc__ = model.description
 
     def __call__(self, *args, **kwargs) -> Any:
@@ -75,14 +78,8 @@ def __call__(self, *args, **kwargs) -> Any:
             TypeError: If the supplied arguments do not match the plan
                 parameter schema.
         """
-        req = TaskRequest(
-            name=self.name,
-            args=list(args),
-            params=kwargs,
-            instrument_session=self._client.instrument_session,
-        )
         try:
-            match self._client.run_task(req):
+            match self._runner(self.name, args, kwargs):
                 case TaskStatus(result=TaskResult(result=res)):
                     return res
                 case TaskStatus(result=TaskError(type=typ, message=msg)):

src/blueapi/client/rest.py

@@ -16,7 +16,6 @@
 from websockets.sync.client import connect
 
 from blueapi import __version__
-from blueapi.client.protocols import ClientObjectRef
 from blueapi.config import RestConfig
 from blueapi.core.bluesky_types import DataEvent
 from blueapi.service.authentication import JWTAuth, SessionManager
@@ -201,6 +200,11 @@ def _response_json(response: requests.Response) -> Any:
         ) from exc
 
 
+class ClientObjectRef:
+    name: str
+    model: DeviceModel
+
+
 class BlueapiRestClient:
     _config: RestConfig
     session_manager: SessionManager | None

src/blueapi/stubs/templates/device_cache_template.pyi.jinja

@@ -0,0 +1,29 @@
+from blueapi.client.rest import BlueapiRestClient
+from blueapi.service.model import DeviceModel
+
+{#-
+ This file is based on the devices_cache.py file in
+ blueapi/client/device_cache.py and should be kept in sync with changes there.
+#}
+
+# This file is auto-generated for a live server and should not be modified directly
+
+class DeviceRef:
+    model: DeviceModel
+    _cache: DeviceCache
+    def __new__(cls, name: str, cache: DeviceCache, model: DeviceModel): ...
+    def __getattr__(self, name) -> DeviceRef: ...
+    def __repr__(self) -> str: ...
+
+class DeviceCache:
+    def __init__(self, rest: BlueapiRestClient) -> None: ...
+    def __getitem__(self, name: str) -> DeviceRef: ...
+    def __iter__(self):  # -> Iterator[DeviceRef]:
+        ...
+    def __repr__(self) -> str: ...
+
+### Generated
+    {%- for item in devices %}
+    {{ item.name }}: DeviceRef
+    {%- endfor %}
+### End