add connect for lazy devices

src/blueapi/core/context.py

@@ -28,8 +28,10 @@
 
 LOGGER = logging.getLogger(__name__)
 
+DevicesDict = dict[str, Device]
 
-def bisect_devices_dict(input_dict: dict[str, Device]) -> tuple[dict]:
+
+def bisect_devices_dict(input_dict: DevicesDict) -> tuple[DevicesDict, DevicesDict]:
     lazy_dict = {}
     non_lazy_dict = {}
 
@@ -56,7 +58,6 @@ class BlueskyContext:
     )
     plans: dict[str, Plan] = field(default_factory=dict)
     devices: dict[str, Device] = field(default_factory=dict)
-    # todo add some format to keep the lazy vs non-lazy devices
 
     plan_functions: dict[str, PlanGenerator] = field(default_factory=dict)
 
@@ -118,15 +119,14 @@ def plan_2(...) -> MsgGenerator:
     def with_device_module(self, module: ModuleType) -> None:
         self.with_dodal_module(module)
 
+    def get_lazy_devices(self) -> DevicesDict:
+        return {k: v for k, v in self.devices.items() if v.lazy}
+
     def with_dodal_module(self, module: ModuleType, **kwargs) -> None:
         devices, exceptions = make_all_devices(module, **kwargs)
-        # todo modify here
-        # factories = get_device_factories(module)
 
         # for non-lazy devices, we instantiate them
-        early_devices, lazy_devices = bisect_devices_dict(devices)
-        # todo for lazy devices we add to the context to do when the plan is run
-        # might need to find that inside the worker
+        early_devices, _ = bisect_devices_dict(devices)
 
         for device in early_devices.values():
             self.register_device(device)

src/blueapi/worker/task.py

@@ -2,6 +2,7 @@
 from collections.abc import Mapping
 from typing import Any
 
+from ophyd_async.plan_stubs import ensure_connected
 from pydantic import BaseModel, Field
 
 from blueapi.core import BlueskyContext
@@ -26,8 +27,8 @@ def prepare_params(self, ctx: BlueskyContext) -> BaseModel:
     def do_task(self, ctx: BlueskyContext) -> None:
         LOGGER.info(f"Asked to run plan {self.name} with {self.params}")
 
-        # todo here call ensure_connected on alpl the devices marked in the context as needing that
-        # ensure_connected(ctx.lazy_devices.values())
+        lazy_devices = ctx.get_lazy_devices()
+        ensure_connected(devices=lazy_devices.values())
         func = ctx.plan_functions[self.name]
         prepared_params = self.prepare_params(ctx)
         ctx.run_engine(func(**prepared_params.dict()))