Generate cache type stubs to help edit user scripts

pyproject.toml

@@ -37,6 +37,7 @@ dependencies = [
     "pyjwt[crypto]",
     "tomlkit",
     "graypy>=2.1.0",
+    "jinja2>=3.1.6",
 ]
 dynamic = ["version"]
 license.file = "LICENSE"

src/blueapi/cli/cli.py

@@ -39,6 +39,7 @@
 from blueapi.service.model import DeviceResponse, PlanResponse, SourceInfo, TaskRequest
 from blueapi.worker import ProgressEvent, WorkerEvent
 
+from . import stubgen
 from .scratch import setup_scratch
 from .updates import CliEventRenderer
 
@@ -152,6 +153,23 @@ 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):
+    """
+    Generate a type-stubs project for blueapi for the currently running server.
+    This enables users using blueapi as a library to benefit from type checking
+    and linting when writing scripts against the BlueapiClient.
+    """
+    click.echo(f"Writing stubs to {target}")
+
+    config: ApplicationConfig = obj["config"]
+    bc = BlueapiClient.from_config(config)
+
+    stubgen.generate_stubs(Path(target), list(bc.plans), list(bc.devices))
+
+
 @main.group()
 @click.option(
     "-o",

src/blueapi/cli/stubgen.py

@@ -0,0 +1,117 @@
+import logging
+from dataclasses import dataclass
+from inspect import cleandoc
+from pathlib import Path
+from textwrap import dedent
+from typing import Self, TextIO
+
+from jinja2 import Environment, PackageLoader
+
+from blueapi.client.cache import DeviceRef, Plan
+from blueapi.core import context
+from blueapi.core.bluesky_types import BLUESKY_PROTOCOLS
+
+log = logging.getLogger(__name__)
+
+
+@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":
+            return f"list[{_type_string(spec.get('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)
+
+    stub_file = client_dir / "cache.pyi"
+    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 file to %s", stub_file)
+    with open(stub_file, "w") as out:
+        render_stub_file(out, plans, 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, plan_models: list[Plan], devices: list[DeviceRef]
+):
+    plans = [PlanSpec.from_plan(p) for p in plan_models]
+
+    env = Environment(loader=PackageLoader("blueapi", package_path="stubs/templates"))
+    env.filters["docstring"] = _docstring
+    tmpl = env.get_template("cache_template.pyi.jinja")
+    stub_file.write(tmpl.render(plans=plans, devices=devices))

src/blueapi/client/cache.py

@@ -0,0 +1,177 @@
+import logging
+from collections.abc import Callable
+from itertools import chain
+from typing import Any
+
+from blueapi.client.rest import BlueapiRestClient
+from blueapi.service.model import DeviceModel, PlanModel
+from blueapi.worker.event import WorkerEvent
+
+log = logging.getLogger(__name__)
+
+
+# This file should be kept in sync with the type stub template in stubs/templates
+
+
+PlanRunner = Callable[[str, dict[str, Any]], WorkerEvent]
+
+
+class PlanCache:
+    """
+    Cache of plans available on the server
+    """
+
+    def __init__(self, runner: PlanRunner, plans: list[PlanModel]):
+        self._cache = {model.name: Plan(model=model, runner=runner) for model in plans}
+        for name, plan in self._cache.items():
+            if name.startswith("_"):
+                continue
+            setattr(self, name, plan)
+
+    def __getitem__(self, name: str) -> "Plan":
+        return self._cache[name]
+
+    def __getattr__(self, name: str) -> "Plan":
+        raise AttributeError(f"No plan named '{name}' available")
+
+    def __iter__(self):
+        return iter(self._cache.values())
+
+    def __repr__(self) -> str:
+        return f"PlanCache({len(self._cache)} plans)"
+
+
+class Plan:
+    """
+    An interface to a plan on the blueapi server
+
+    This allows remote plans to be called (mostly) as if they were local
+    methods when writing user scripts.
+
+    If you are seeing this help while using blueapi as a library, generating
+    type stubs may be helpful for type checking and plan discovery, eg
+
+        blueapi generate-stubs /tmp/blueapi-stubs
+        uv add --editable /tmp/blueapi-stubs
+
+    """
+
+    model: PlanModel
+
+    def __init__(self, model: PlanModel, runner: PlanRunner):
+        self.model = model
+        self._runner = runner
+        self.__doc__ = model.description
+
+    def __call__(self, *args, **kwargs) -> WorkerEvent:
+        """
+        Run the plan on the server mapping the given args into the required parameters
+        """
+        return self._runner(self.name, self._build_args(*args, **kwargs))
+
+    @property
+    def name(self) -> str:
+        return self.model.name
+
+    @property
+    def help_text(self) -> str:
+        return self.model.description or f"Plan {self!r}"
+
+    @property
+    def properties(self) -> set[str]:
+        return self.model.parameter_schema.get("properties", {}).keys()
+
+    @property
+    def required(self) -> list[str]:
+        return self.model.parameter_schema.get("required", [])
+
+    def _build_args(self, *args, **kwargs):
+        log.info(
+            "Building args for %s, using %s and %s",
+            "[" + ",".join(self.properties) + "]",
+            args,
+            kwargs,
+        )
+
+        if len(args) > len(self.properties):
+            raise TypeError(f"{self.name} got too many arguments")
+        if extra := {k for k in kwargs if k not in self.properties}:
+            raise TypeError(f"{self.name} got unexpected arguments: {extra}")
+
+        params = {}
+        # Initially fill parameters using positional args assuming the order
+        # from the parameter_schema
+        for req, arg in zip(self.properties, args, strict=False):
+            params[req] = arg
+
+        # Then append any values given via kwargs
+        for key, value in kwargs.items():
+            # If we've already assumed a positional arg was this value, bail out
+            if key in params:
+                raise TypeError(f"{self.name} got multiple values for {key}")
+            params[key] = value
+
+        if missing := {k for k in self.required if k not in params}:
+            raise TypeError(f"Missing argument(s) for {missing}")
+        return params
+
+    def __repr__(self):
+        opts = [p for p in self.properties if p not in self.required]
+        params = ", ".join(chain(self.required, (f"{opt}=None" for opt in opts)))
+        return f"{self.name}({params})"
+
+
+class DeviceCache:
+    def __init__(self, rest: BlueapiRestClient):
+        self._rest = rest
+        self._cache = {
+            model.name: DeviceRef(name=model.name, cache=self, model=model)
+            for model in rest.get_devices().devices
+        }
+        for name, device in self._cache.items():
+            if name.startswith("_"):
+                continue
+            setattr(self, name, device)
+
+    def __getitem__(self, name: str) -> "DeviceRef":
+        if dev := self._cache.get(name):
+            return dev
+        try:
+            model = self._rest.get_device(name)
+            device = DeviceRef(name=name, cache=self, model=model)
+            self._cache[name] = device
+            setattr(self, model.name, device)
+            return device
+        except KeyError:
+            pass
+        raise AttributeError(f"No device named '{name}' available")
+
+    def __getattr__(self, name: str) -> "DeviceRef":
+        if name.startswith("_"):
+            return super().__getattribute__(name)
+        return self[name]
+
+    def __iter__(self):
+        return iter(self._cache.values())
+
+    def __repr__(self) -> str:
+        return f"DeviceCache({len(self._cache)} devices)"
+
+
+class DeviceRef(str):
+    model: DeviceModel
+    _cache: DeviceCache
+
+    def __new__(cls, name: str, cache: DeviceCache, model: DeviceModel):
+        instance = super().__new__(cls, name)
+        instance.model = model
+        instance._cache = cache
+        return instance
+
+    def __getattr__(self, name) -> "DeviceRef":
+        if name.startswith("_"):
+            raise AttributeError(f"No child device named {name}")
+        return self._cache[f"{self}.{name}"]
+
+    def __repr__(self):
+        return f"Device({self})"