diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index 076b43f5e..379b211da 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -29,6 +29,20 @@ LOGGER = logging.getLogger(__name__) +def bisect_devices_dict(input_dict: dict[str, Device]) -> tuple[dict]: + lazy_dict = {} + non_lazy_dict = {} + + for key, value in input_dict.items(): + # Assuming 'lazy' is an attribute or a condition you can check + if hasattr(value, "lazy") and value.lazy: + lazy_dict[key] = value + else: + non_lazy_dict[key] = value + + return (non_lazy_dict, lazy_dict) + + @dataclass class BlueskyContext: """ @@ -43,6 +57,7 @@ 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) _reference_cache: dict[type, type] = field(default_factory=dict) @@ -98,7 +113,7 @@ def plan_2(...) -> MsgGenerator: for obj in load_module_all(module): if is_bluesky_plan_generator(obj): - self.plan(obj) + self.register_plan(obj) def with_device_module(self, module: ModuleType) -> None: self.with_dodal_module(module) @@ -109,10 +124,12 @@ def with_dodal_module(self, module: ModuleType, **kwargs) -> None: # factories = get_device_factories(module) # for non-lazy devices, we instantiate them - # for lazy devices we add to the context to do when the plan is run + 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 - for device in devices.values(): - self.device(device) + for device in early_devices.values(): + self.register_device(device) # If exceptions have occurred, we log them but we do not make blueapi # fall over @@ -122,7 +139,7 @@ def with_dodal_module(self, module: ModuleType, **kwargs) -> None: ) LOGGER.exception(NotConnected(exceptions)) - def plan(self, plan: PlanGenerator) -> PlanGenerator: + def register_plan(self, plan: PlanGenerator) -> PlanGenerator: """ Register the argument as a plan in the context. Can be used as a decorator e.g. @ctx.plan @@ -150,7 +167,7 @@ def my_plan(a: int, b: str): self.plan_functions[plan.__name__] = plan return plan - def device(self, device: Device, name: str | None = None) -> None: + def register_device(self, device: Device, name: str | None = None) -> None: """ Register an device in the context. The device needs to be registered with a name. If the device is Readable, Movable or Flyable it has a `name` diff --git a/tests/core/test_context.py b/tests/core/test_context.py index a1f6bdb1a..3e746b963 100644 --- a/tests/core/test_context.py +++ b/tests/core/test_context.py @@ -96,8 +96,8 @@ def empty_context() -> BlueskyContext: @pytest.fixture def devicey_context(sim_motor: SynAxis, sim_detector: SynGauss) -> BlueskyContext: ctx = BlueskyContext() - ctx.device(sim_motor) - ctx.device(sim_detector) + ctx.register_device(sim_motor) + ctx.register_device(sim_detector) return ctx @@ -117,7 +117,7 @@ def some_configurable() -> SomeConfigurable: @pytest.mark.parametrize("plan", [has_no_params, has_one_param, has_some_params]) def test_add_plan(empty_context: BlueskyContext, plan: PlanGenerator) -> None: - empty_context.plan(plan) + empty_context.register_plan(plan) assert plan.__name__ in empty_context.plans @@ -127,7 +127,7 @@ def test_generated_schema( def demo_plan(foo: int, mov: Movable) -> MsgGenerator: # type: ignore ... - empty_context.plan(demo_plan) + empty_context.register_plan(demo_plan) schema = empty_context.plans["demo_plan"].model.schema() assert schema["properties"] == { "foo": {"title": "Foo", "type": "integer"}, @@ -140,7 +140,7 @@ def demo_plan(foo: int, mov: Movable) -> MsgGenerator: # type: ignore ) def test_add_invalid_plan(empty_context: BlueskyContext, plan: PlanGenerator) -> None: with pytest.raises(ValueError): - empty_context.plan(plan) + empty_context.register_plan(plan) def test_add_plan_from_module(empty_context: BlueskyContext) -> None: @@ -151,14 +151,14 @@ def test_add_plan_from_module(empty_context: BlueskyContext) -> None: def test_add_named_device(empty_context: BlueskyContext, sim_motor: SynAxis) -> None: - empty_context.device(sim_motor) + empty_context.register_device(sim_motor) assert empty_context.devices[SIM_MOTOR_NAME] is sim_motor def test_add_nameless_device( empty_context: BlueskyContext, some_configurable: SomeConfigurable ) -> None: - empty_context.device(some_configurable, "conf") + empty_context.register_device(some_configurable, "conf") assert empty_context.devices["conf"] is some_configurable @@ -167,13 +167,13 @@ def test_add_nameless_device_without_override( some_configurable: SomeConfigurable, ) -> None: with pytest.raises(KeyError): - empty_context.device(some_configurable) + empty_context.register_device(some_configurable) def test_override_device_name( empty_context: BlueskyContext, sim_motor: SynAxis ) -> None: - empty_context.device(sim_motor, "foo") + empty_context.register_device(sim_motor, "foo") assert empty_context.devices["foo"] is sim_motor @@ -246,12 +246,12 @@ def test_lookup_non_device(devicey_context: BlueskyContext) -> None: def test_add_non_plan(empty_context: BlueskyContext) -> None: with pytest.raises(TypeError): - empty_context.plan("not a plan") # type: ignore + empty_context.register_plan("not a plan") # type: ignore def test_add_non_device(empty_context: BlueskyContext) -> None: with pytest.raises(TypeError): - empty_context.device("not a device") # type: ignore + empty_context.register_device("not a device") # type: ignore def test_add_devices_and_plans_from_modules_with_config( @@ -361,8 +361,8 @@ def test_str_default( empty_context: BlueskyContext, sim_motor: SynAxis, alt_motor: SynAxis ): movable_ref = empty_context._reference(Movable) - empty_context.device(sim_motor) - empty_context.plan(has_default_reference) + empty_context.register_device(sim_motor) + empty_context.register_plan(has_default_reference) spec = empty_context._type_spec_for_function(has_default_reference) assert spec["m"][0] is movable_ref @@ -371,7 +371,7 @@ def test_str_default( assert has_default_reference.__name__ in empty_context.plans model = empty_context.plans[has_default_reference.__name__].model assert parse_obj_as(model, {}).m is sim_motor # type: ignore - empty_context.device(alt_motor) + empty_context.register_device(alt_motor) assert parse_obj_as(model, {"m": ALT_MOTOR_NAME}).m is alt_motor # type: ignore @@ -379,8 +379,8 @@ def test_nested_str_default( empty_context: BlueskyContext, sim_motor: SynAxis, alt_motor: SynAxis ): movable_ref = empty_context._reference(Movable) - empty_context.device(sim_motor) - empty_context.plan(has_default_nested_reference) + empty_context.register_device(sim_motor) + empty_context.register_plan(has_default_nested_reference) spec = empty_context._type_spec_for_function(has_default_nested_reference) assert spec["m"][0] == list[movable_ref] # type: ignore @@ -389,7 +389,7 @@ def test_nested_str_default( assert has_default_nested_reference.__name__ in empty_context.plans model = empty_context.plans[has_default_nested_reference.__name__].model assert parse_obj_as(model, {}).m == [sim_motor] # type: ignore - empty_context.device(alt_motor) + empty_context.register_device(alt_motor) assert parse_obj_as(model, {"m": [ALT_MOTOR_NAME]}).m == [alt_motor] # type: ignore @@ -398,6 +398,6 @@ def a_plan(foo_bar: int, baz: str) -> MsgGenerator: if False: yield - empty_context.plan(a_plan) + empty_context.register_plan(a_plan) with pytest.raises(ValidationError): empty_context.plans[a_plan.__name__].model(fooBar=1, baz="test") diff --git a/tests/worker/test_reworker.py b/tests/worker/test_reworker.py index 8e5fc217b..d488ad1bc 100644 --- a/tests/worker/test_reworker.py +++ b/tests/worker/test_reworker.py @@ -65,8 +65,8 @@ def context(fake_device: FakeDevice) -> BlueskyContext: ctx_config.sources.append( Source(kind=SourceKind.DEVICE_FUNCTIONS, module="devices") ) - ctx.plan(failing_plan) - ctx.device(fake_device) + ctx.register_plan(failing_plan) + ctx.register_device(fake_device) ctx.with_config(ctx_config) return ctx