Generate cache type stubs to help edit user scripts #1338

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft

tpoliaw wants to merge 35 commits into user-client from cache-stubs

pyproject.toml

-Original file line number
+Diff line change
@@ Expand Up / @@ -37,6 +37,7 @@ dependencies = [ @@
         "pyjwt[crypto]",
         "tomlkit",
         "graypy>=2.1.0",
+        "jinja2>=3.1.6",
     ]
     dynamic = ["version"]
     license.file = "LICENSE"
@@ Expand Down @@

src/blueapi/cli/cli.py

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -17,7 +17,6 @@
  
    from click.exceptions import ClickException

    from observability_utils.tracing import setup_tracing

    from pydantic import ValidationError

    from requests.exceptions import ConnectionError

    from blueapi import __version__, config

    from blueapi.cli.format import OutputFormat

    @@ -26,6 +25,7 @@
  
    from blueapi.client.rest import (

        BlueskyRemoteControlError,

        InvalidParametersError,

        ServiceUnavailableError,

        UnauthorisedAccessError,

        UnknownPlanError,

    )

    @@ -36,9 +36,10 @@
  
    from blueapi.core import OTLP_EXPORT_ENABLED, DataEvent

    from blueapi.log import set_up_logging

    from blueapi.service.authentication import SessionCacheManager, SessionManager

    from blueapi.service.model import SourceInfo, TaskRequest

    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",

    @@ -183,7 +201,7 @@ def check_connection(func: Callable[P, T]) -> Callable[P, T]:
  
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:

            try:

                return func(*args, **kwargs)

            except ConnectionError as ce:

            except ServiceUnavailableError as ce:

                raise ClickException(

                    "Failed to establish connection to blueapi server."

                ) from ce

    @@ -204,7 +222,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
  
    def get_plans(obj: dict) -> None:

        """Get a list of plans available for the worker to use"""

        client: BlueapiClient = obj["client"]

        obj["fmt"].display(client.get_plans())

        obj["fmt"].display(PlanResponse(plans=[p.model for p in client.plans]))

    @controller.command(name="devices")

    @@ -213,7 +231,7 @@ def get_plans(obj: dict) -> None:
  
    def get_devices(obj: dict) -> None:

        """Get a list of devices available for the worker to use"""

        client: BlueapiClient = obj["client"]

        obj["fmt"].display(client.get_devices())

        obj["fmt"].display(DeviceResponse(devices=[dev.model for dev in client.devices]))

    @controller.command(name="listen")

    @@ -345,7 +363,7 @@ def get_state(obj: dict) -> None:
  
        """Print the current state of the worker"""

        client: BlueapiClient = obj["client"]

        print(client.get_state().name)

        print(client.state.name)

    @controller.command(name="pause")

    @@ -428,7 +446,7 @@ def env(
  
            status = client.reload_environment(timeout=timeout)

            print("Environment is initialized")

        else:

            status = client.get_environment()

            status = client.environment

        print(status)

    @@ -470,14 +488,13 @@ def login(obj: dict) -> None:
  
            print("Logged in")

        except Exception:

            client = BlueapiClient.from_config(config)

            oidc_config = client.get_oidc_config()

            if oidc_config is None:

            if oidc := client.oidc_config:

                auth = SessionManager(

                    oidc, cache_manager=SessionCacheManager(config.auth_token_path)

                )

                auth.start_device_flow()

            else:

                print("Server is not configured to use authentication!")

                return

            auth = SessionManager(

                oidc_config, cache_manager=SessionCacheManager(config.auth_token_path)

            )

            auth.start_device_flow()

    @main.command(name="logout")

src/blueapi/cli/format.py

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -12,7 +12,9 @@
  
    from blueapi.core.bluesky_types import DataEvent

    from blueapi.service.model import (

        DeviceModel,

        DeviceResponse,

        PlanModel,

        PlanResponse,

        PythonEnvironmentResponse,

        SourceInfo,

    @@ -54,17 +56,21 @@ def display_full(obj: Any, stream: Stream):
  
        match obj:

            case PlanResponse(plans=plans):

                for plan in plans:

                    print(plan.name)

                    if desc := plan.description:

                        print(indent(dedent(desc).strip(), "    "))

                    if schema := plan.parameter_schema:

                        print("    Schema")

                        print(indent(json.dumps(schema, indent=2), "        "))

                    display_full(plan, stream)

            case PlanModel(name=name, description=desc, parameter_schema=schema):

                print(name)

                if desc:

                    print(indent(dedent(desc).strip(), "    "))

                if schema:

                    print("    Schema")

                    print(indent(json.dumps(schema, indent=2), "        "))

            case DeviceResponse(devices=devices):

                for dev in devices:

                    print(dev.name)

                    for proto in dev.protocols:

                        print(f"    {proto}")

                    display_full(dev, stream)

            case DeviceModel(name=name, protocols=protocols):

                print(name)

                for proto in protocols:

                    print(f"    {proto}")

            case DataEvent(name=name, doc=doc):

                print(f"{name.title()}:{fmt_dict(doc)}")

            case WorkerEvent(state=st, task_status=task):

    @@ -100,11 +106,13 @@ def display_json(obj: Any, stream: Stream):
  
        print = partial(builtins.print, file=stream)

        match obj:

            case PlanResponse(plans=plans):

                print(json.dumps([p.model_dump() for p in plans], indent=2))

                display_json(plans, stream)

            case DeviceResponse(devices=devices):

                print(json.dumps([d.model_dump() for d in devices], indent=2))

                display_json(devices, stream)

            case BaseModel():

                print(json.dumps(obj.model_dump()))

            case list():

                print(json.dumps([it.model_dump() for it in obj], indent=2))

            case _:

                print(json.dumps(obj))

    @@ -114,26 +122,30 @@ def display_compact(obj: Any, stream: Stream):
  
        match obj:

            case PlanResponse(plans=plans):

                for plan in plans:

                    print(plan.name)

                    if desc := plan.description:

                        print(indent(dedent(desc.split("\n\n")[0].strip("\n")), "    "))

                    if schema := plan.parameter_schema:

                        print("    Args")

                        for arg, spec in schema.get("properties", {}).items():

                            req = arg in schema.get("required", {})

                            print(f"      {arg}={_describe_type(spec, req)}")

                    display_compact(plan, stream)

            case PlanModel(name=name, description=desc, parameter_schema=schema):

                print(name)

                if desc:

                    print(indent(dedent(desc.split("\n\n")[0].strip("\n")), "    "))

                if schema:

                    print("    Args")

                    for arg, spec in schema.get("properties", {}).items():

                        req = arg in schema.get("required", {})

                        print(f"      {arg}={_describe_type(spec, req)}")

            case DeviceResponse(devices=devices):

                for dev in devices:

                    print(dev.name)

                    print(

                        indent(

                            textwrap.fill(

                                ", ".join(str(proto) for proto in dev.protocols),

                                80,

                            ),

                            "    ",

                        )

                    display_compact(dev, stream)

            case DeviceModel(name=name, protocols=protocols):

                print(name)

                print(

                    indent(

                        textwrap.fill(

                            ", ".join(str(proto) for proto in protocols),

                            80,

                        ),

                        "    ",

                    )

                )

            case DataEvent(name=name):

                print(f"Data Event: {name}")

            case WorkerEvent(state=state):

src/blueapi/cli/stubgen.py

-Original file line number
+Diff line change
@@ -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))

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Generate cache type stubs to help edit user scripts #1338

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Generate cache type stubs to help edit user scripts #1338

Are you sure you want to change the base?

Uh oh!

Generate cache type stubs to help edit user scripts #1338

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!