diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index 379b211da..fac18b429 100644 --- a/src/blueapi/core/context.py +++ b/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) diff --git a/src/blueapi/worker/task.py b/src/blueapi/worker/task.py index 47708e3f4..37f1b6d40 100644 --- a/src/blueapi/worker/task.py +++ b/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()))