Naive generic handling

DiamondLightSource · Mar 3, 2025 · 55ba471 · 55ba471
1 parent 60a74ab
commit 55ba471
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 7 deletions.
diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py
@@ -36,6 +36,12 @@
 LOGGER = logging.getLogger(__name__)
 
 
+def _type_name(target: type):
+    if (module := target.__module__) != "builtins":
+        return f"{module}.{target.__name__ or target.__qualname__}"
+    return f"{target.__name__ or target.__qualname__}"
+
+
 @dataclass
 class BlueskyContext:
     """
@@ -232,7 +238,11 @@ def __get_pydantic_json_schema__(
                 ) -> JsonSchemaValue:
                     json_schema = handler(core_schema)
                     json_schema = handler.resolve_ref_schema(json_schema)
-                    json_schema["type"] = f"{target.__module__}.{target.__qualname__}"
+                    json_schema["type"] = (
+                        _type_name(target) + f"{[_type_name(arg) for arg in cls.args]}"
+                        if cls.args
+                        else ""
+                    )
                     return json_schema
 
             self._reference_cache[target] = Reference

diff --git a/tests/system_tests/plans.json b/tests/system_tests/plans.json
@@ -2,7 +2,7 @@
     "plans": [
         {
             "name": "stp_snapshot",
-            "description": "\n    Moves devices for pressure and temperature (defaults fetched from the context)\n    and captures a single frame from a collection of devices\n\n    Args:\n        detectors (List[Readable]): A list of devices to read while the sample is at STP\n        temperature (Optional[Movable]): A device controlling temperature of the sample,\n            defaults to fetching a device name \"sample_temperature\" from the context\n        pressure (Optional[Movable]): A device controlling pressure on the sample,\n            defaults to fetching a device name \"sample_pressure\" from the context\n    Returns:\n        MsgGenerator: Plan\n    Yields:\n        Iterator[MsgGenerator]: Bluesky messages\n    ",
+            "description": "\n    Moves devices for pressure and temperature (defaults fetched from the context)\n    and captures a single frame from a collection of devices\n\n    Args:\n        detectors (List[Readable]): A list of devices to read while the sample is at STP\n        temperature (Optional[[Movable[float]]): A device controlling temperature of\n            the sample, defaults to fetching a device named \"sample_temperature\" from\n            the context\n        pressure (Optional[[Movable[float]]): A device controlling pressure on the\n            sample, defaults to fetching a device named \"sample_pressure\" from the\n            context\n    Returns:\n        MsgGenerator: Plan\n    Yields:\n        Iterator[MsgGenerator]: Bluesky messages\n    ",
             "schema": {
                 "additionalProperties": false,
                 "properties": {
@@ -15,11 +15,11 @@
                     },
                     "temperature": {
                         "title": "Temperature",
-                        "type": "bluesky.protocols.Movable"
+                        "type": "bluesky.protocols.Movable['float']"
                     },
                     "pressure": {
                         "title": "Pressure",
-                        "type": "bluesky.protocols.Movable"
+                        "type": "bluesky.protocols.Movable['float']"
                     }
                 },
                 "required": [
@@ -936,13 +936,13 @@
         },
         {
             "name": "set_absolute",
-            "description": "\n    Set a device, wrapper for `bp.abs_set`.\n\n    Args:\n        movable (Movable): The device to set\n        value (T): The new value\n        group (Group | None, optional): The message group to associate with the\n                                           setting, for sequencing. Defaults to None.\n        wait (bool, optional): The group should wait until all setting is complete\n                               (e.g. a motor has finished moving). Defaults to False.\n\n    Returns:\n        MsgGenerator: Plan\n\n    Yields:\n        Iterator[MsgGenerator]: Bluesky messages\n    ",
+            "description": "\n    Set a device, wrapper for `bp.abs_set`.\n\n    Args:\n        movable (Movable[T]): The device to set\n        value (T): The new value\n        group (Group | None, optional): The message group to associate with the\n                                           setting, for sequencing. Defaults to None.\n        wait (bool, optional): The group should wait until all setting is complete\n                               (e.g. a motor has finished moving). Defaults to False.\n\n    Returns:\n        MsgGenerator: Plan\n\n    Yields:\n        Iterator[MsgGenerator]: Bluesky messages\n    ",
             "schema": {
                 "additionalProperties": false,
                 "properties": {
                     "movable": {
                         "title": "Movable",
-                        "type": "bluesky.protocols.Movable"
+                        "type": "bluesky.protocols.Movable['dodal.plan_stubs.wrapped.T']"
                     },
                     "value": {
                         "title": "Value"
@@ -979,7 +979,7 @@
                 "properties": {
                     "movable": {
                         "title": "Movable",
-                        "type": "bluesky.protocols.Movable"
+                        "type": "bluesky.protocols.Movable['dodal.plan_stubs.wrapped.T']"
                     },
                     "value": {
                         "title": "Value"