From 43a3090a5b1907435c42c3955e273af7ed952763 Mon Sep 17 00:00:00 2001 From: Callum Forrester Date: Fri, 5 May 2023 13:23:36 +0100 Subject: [PATCH 1/6] Remove public fields of context and replace with methods --- src/blueapi/core/bluesky_types.py | 2 +- src/blueapi/core/context.py | 46 +++++++++++++++++++++++-------- src/blueapi/service/app.py | 4 +-- src/blueapi/worker/task.py | 13 +++++---- 4 files changed, 46 insertions(+), 19 deletions(-) diff --git a/src/blueapi/core/bluesky_types.py b/src/blueapi/core/bluesky_types.py index a5cfa8467..78ed9db20 100644 --- a/src/blueapi/core/bluesky_types.py +++ b/src/blueapi/core/bluesky_types.py @@ -84,7 +84,7 @@ def is_bluesky_plan_generator(func: PlanGenerator) -> bool: class Plan(BlueapiBaseModel): """ - A plan that can be run + Metadata about a plan that can be run """ name: str = Field(description="Referenceable name of the plan") diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index 9d28d4f6c..c69b00a2c 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -8,6 +8,7 @@ Any, Callable, Dict, + Iterable, List, Optional, Tuple, @@ -37,21 +38,44 @@ LOGGER = logging.getLogger(__name__) -@dataclass class BlueskyContext: """ Context for building a Bluesky application """ - run_engine: RunEngine = field( - default_factory=lambda: RunEngine(context_managers=[]) - ) - plans: Dict[str, Plan] = field(default_factory=dict) - devices: Dict[str, Device] = field(default_factory=dict) - plan_functions: Dict[str, PlanGenerator] = field(default_factory=dict) + _run_engine: RunEngine + _plans: Dict[str, Plan] + _devices: Dict[str, Device] + _plan_functions: Dict[str, PlanGenerator] _reference_cache: Dict[Type, Type] = field(default_factory=dict) + def __init__(self, run_engine: Optional[RunEngine] = None) -> None: + if run_engine is None: + run_engine = RunEngine(context_managers=[]) + + self._run_engine = run_engine + self._devices = {} + self._plans = {} + self._plan_functions = {} + self._reference_cache = {} + + @property + def run_engine(self) -> RunEngine: + return self._run_engine + + def find_plan_function(self, name: str) -> Optional[PlanGenerator]: + return self._plan_functions.get(name) + + def find_plan_metadata(self, name: str) -> Optional[Plan]: + return self._plans.get(name) + + def all_devices(self) -> Iterable[Device]: + return self._devices.values() + + def all_plan_metadata(self) -> Iterable[Plan]: + return self._plans.values() + def find_device(self, addr: Union[str, List[str]]) -> Optional[Device]: """ Find a device in this context, allows for recursive search. @@ -68,7 +92,7 @@ def find_device(self, addr: Union[str, List[str]]) -> Optional[Device]: list_addr = list(addr.split(".")) return self.find_device(list_addr) else: - return find_component(self.devices, addr) + return find_component(self._devices, addr) def with_startup_script(self, path: Union[Path, str]) -> None: mod = import_module(str(path)) @@ -129,8 +153,8 @@ def my_plan(a: int, b: str): __config__=BlueapiPlanModelConfig, **self._type_spec_for_function(plan), ) - self.plans[plan.__name__] = Plan(name=plan.__name__, model=model) - self.plan_functions[plan.__name__] = plan + self._plans[plan.__name__] = Plan(name=plan.__name__, model=model) + self._plan_functions[plan.__name__] = plan return plan def device(self, device: Device, name: Optional[str] = None) -> None: @@ -158,7 +182,7 @@ def device(self, device: Device, name: Optional[str] = None) -> None: else: raise KeyError(f"Must supply a name for this device: {device}") - self.devices[name] = device + self._devices[name] = device def _reference(self, target: Type) -> Type: """ diff --git a/src/blueapi/service/app.py b/src/blueapi/service/app.py index 3a2305407..9e3b4937d 100644 --- a/src/blueapi/service/app.py +++ b/src/blueapi/service/app.py @@ -81,7 +81,7 @@ def _on_run_request(self, message_context: MessageContext, task: RunPlan) -> Non self._template.send(reply_queue, response) def _get_plans(self, message_context: MessageContext, message: PlanRequest) -> None: - plans = list(map(PlanModel.from_plan, self._ctx.plans.values())) + plans = list(map(PlanModel.from_plan, self._ctx.all_plan_metadata())) response = PlanResponse(plans=plans) assert message_context.reply_destination is not None @@ -90,7 +90,7 @@ def _get_plans(self, message_context: MessageContext, message: PlanRequest) -> N def _get_devices( self, message_context: MessageContext, message: DeviceRequest ) -> None: - devices = list(map(DeviceModel.from_device, self._ctx.devices.values())) + devices = list(map(DeviceModel.from_device, self._ctx.all_devices())) response = DeviceResponse(devices=devices) assert message_context.reply_destination is not None diff --git a/src/blueapi/worker/task.py b/src/blueapi/worker/task.py index 2a77a18ba..11bd25736 100644 --- a/src/blueapi/worker/task.py +++ b/src/blueapi/worker/task.py @@ -41,15 +41,18 @@ class RunPlan(Task): def do_task(self, ctx: BlueskyContext) -> None: LOGGER.info(f"Asked to run plan {self.name} with {self.params}") - plan = ctx.plans[self.name] - func = ctx.plan_functions[self.name] - sanitized_params = _lookup_params(ctx, plan, self.params) - plan_generator = func(**sanitized_params.dict()) + metadata = ctx.find_plan_metadata(self.name) + func = ctx.find_plan_function(self.name) + if metadata is None or func is None: + raise KeyError(f"No plan named {self.name}") + sanitized_params = _lookup_params(metadata, self.params).dict() + plan_generator = func(**sanitized_params) ctx.run_engine(plan_generator) def _lookup_params( - ctx: BlueskyContext, plan: Plan, params: Mapping[str, Any] + plan: Plan, + params: Mapping[str, Any], ) -> BaseModel: """ Checks plan parameters against context From ccd38443dd58f3dfe3ad604d4ab42231c0c155c8 Mon Sep 17 00:00:00 2001 From: Callum Forrester Date: Fri, 5 May 2023 13:28:21 +0100 Subject: [PATCH 2/6] Fix tests --- src/blueapi/core/context.py | 3 +++ tests/core/test_context.py | 8 ++++---- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index c69b00a2c..c1e71fe55 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -76,6 +76,9 @@ def all_devices(self) -> Iterable[Device]: def all_plan_metadata(self) -> Iterable[Plan]: return self._plans.values() + def has_plan(self, name: str) -> bool: + return name in self._plans + def find_device(self, addr: Union[str, List[str]]) -> Optional[Device]: """ Find a device in this context, allows for recursive search. diff --git a/tests/core/test_context.py b/tests/core/test_context.py index 059a7473a..5c039755d 100644 --- a/tests/core/test_context.py +++ b/tests/core/test_context.py @@ -98,7 +98,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) - assert plan.__name__ in empty_context.plans + assert empty_context.has_plan(plan.__name__) @pytest.mark.parametrize( @@ -111,14 +111,14 @@ def test_add_invalid_plan(empty_context: BlueskyContext, plan: PlanGenerator) -> def test_add_named_device(empty_context: BlueskyContext, sim_motor: SynAxis) -> None: empty_context.device(sim_motor) - assert empty_context.devices[SIM_MOTOR_NAME] is sim_motor + assert empty_context.find_device(SIM_MOTOR_NAME) is sim_motor def test_add_nameless_device( empty_context: BlueskyContext, some_configurable: SomeConfigurable ) -> None: empty_context.device(some_configurable, "conf") - assert empty_context.devices["conf"] is some_configurable + assert empty_context.find_device("conf") is some_configurable def test_add_nameless_device_without_override( @@ -133,7 +133,7 @@ def test_override_device_name( empty_context: BlueskyContext, sim_motor: SynAxis ) -> None: empty_context.device(sim_motor, "foo") - assert empty_context.devices["foo"] is sim_motor + assert empty_context.find_device("foo") is sim_motor @pytest.mark.parametrize( From 38b8387517fd39988c430b6591416e77081027c1 Mon Sep 17 00:00:00 2001 From: Callum Forrester Date: Fri, 5 May 2023 13:42:50 +0100 Subject: [PATCH 3/6] Tidy docstrings in context --- src/blueapi/core/context.py | 86 ++++++++++++++++++++++++++++++++++--- 1 file changed, 81 insertions(+), 5 deletions(-) diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index c1e71fe55..f70bf77e0 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -62,21 +62,73 @@ def __init__(self, run_engine: Optional[RunEngine] = None) -> None: @property def run_engine(self) -> RunEngine: + """ + RunEngine capable of running the plans in the context against the devices in the context. + + Returns: + RunEngine: A Bluesky RunEngine + """ + return self._run_engine def find_plan_function(self, name: str) -> Optional[PlanGenerator]: + """ + Return a plan function matching the given name + + Args: + name: The name of the plan + + Returns: + Optional[PlanGenerator]: The plan function, None if there is + no plan matching the name in the context. + """ return self._plan_functions.get(name) def find_plan_metadata(self, name: str) -> Optional[Plan]: + """ + Return a metadata object describing a plan in the context + + Args: + name: The name of the plan + + Returns: + Optional[Plan]: The plan metadata, None if there is + no plan matching the name in the context. + """ + return self._plans.get(name) def all_devices(self) -> Iterable[Device]: + """ + Return iterable of all devices in the context. + + Returns: + Iterable[Device]: Iterable of devices + """ + return self._devices.values() def all_plan_metadata(self) -> Iterable[Plan]: + """ + Return iterable of all plan metadata in the context. + + Returns: + Iterable[Device]: Iterable of plan metadata + """ + return self._plans.values() def has_plan(self, name: str) -> bool: + """ + Check if a plan exists in the context. + + Args: + name: The name of the plan + + Returns: + bool: True if the plan is accessible from this context + """ + return name in self._plans def find_device(self, addr: Union[str, List[str]]) -> Optional[Device]: @@ -84,11 +136,11 @@ def find_device(self, addr: Union[str, List[str]]) -> Optional[Device]: Find a device in this context, allows for recursive search. Args: - addr (Union[str, List[str]]): Address of the device, examples: - "motors", "motors.x" + addr: Address of the device, examples: "motors", "motors.x" Returns: - Optional[Device]: _description_ + Optional[Device]: A Bluesky compatible device, None if no device + matching addr exists in the context. """ if isinstance(addr, str): @@ -98,10 +150,27 @@ def find_device(self, addr: Union[str, List[str]]) -> Optional[Device]: return find_component(self._devices, addr) def with_startup_script(self, path: Union[Path, str]) -> None: + """ + Register all plans and devices in the Python file at this path. + + Args: + path: Path to the Python file. Can be a posix file path + or a module path if pointing to a valid Python + module in your environment. + """ + mod = import_module(str(path)) self.with_module(mod) def with_module(self, module: ModuleType) -> None: + """ + Register all plans and devices in this module + + Args: + module: A module, usually retrieved from an + import statement + """ + self.with_plan_module(module) self.with_device_module(module) @@ -122,7 +191,7 @@ def plan_2(...): __all__ = ["plan_1", "plan_2"] Args: - module (ModuleType): Module to pass in + module: Module to check for plans """ for obj in load_module_all(module): @@ -130,6 +199,13 @@ def plan_2(...): self.plan(obj) def with_device_module(self, module: ModuleType) -> None: + """ + Register all devices in the module supplied + + Args: + module: Module to check for devices + """ + for obj in load_module_all(module): if is_bluesky_compatible_device(obj): self.device(obj) @@ -164,7 +240,7 @@ def device(self, device: Device, name: Optional[str] = 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` - attribbute which can be used. The attribute can be overrideen with the + attribute which can be used. The attribute can be overridden with the `name` parameter here. If the device conforms to a different protocol then the parameter must be used to name it. From f6ea454c116732dba89ad37c4e135d9ed729b5e9 Mon Sep 17 00:00:00 2001 From: Callum Forrester Date: Fri, 5 May 2023 13:43:40 +0100 Subject: [PATCH 4/6] Fix flake8 errors --- src/blueapi/core/context.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index f70bf77e0..7eb6f9320 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -1,5 +1,5 @@ import logging -from dataclasses import dataclass, field +from dataclasses import field from importlib import import_module from inspect import Parameter, signature from pathlib import Path @@ -47,8 +47,7 @@ class BlueskyContext: _plans: Dict[str, Plan] _devices: Dict[str, Device] _plan_functions: Dict[str, PlanGenerator] - - _reference_cache: Dict[Type, Type] = field(default_factory=dict) + _reference_cache: Dict[Type, Type] def __init__(self, run_engine: Optional[RunEngine] = None) -> None: if run_engine is None: @@ -63,7 +62,8 @@ def __init__(self, run_engine: Optional[RunEngine] = None) -> None: @property def run_engine(self) -> RunEngine: """ - RunEngine capable of running the plans in the context against the devices in the context. + RunEngine capable of running the plans in the context against the devices + in the context. Returns: RunEngine: A Bluesky RunEngine From a6646e70a122a66cfe3614910d4ed4fc384f82aa Mon Sep 17 00:00:00 2001 From: Callum Forrester Date: Fri, 5 May 2023 13:44:39 +0100 Subject: [PATCH 5/6] Remove unused import --- src/blueapi/core/context.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index 7eb6f9320..b10962065 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -1,5 +1,4 @@ import logging -from dataclasses import field from importlib import import_module from inspect import Parameter, signature from pathlib import Path From b8761c87afb428d889a15c0a7c1d03d6bdae64b1 Mon Sep 17 00:00:00 2001 From: Callum Forrester Date: Fri, 5 May 2023 14:30:11 +0100 Subject: [PATCH 6/6] Simplify run engine defaulting --- src/blueapi/core/context.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index b10962065..ab3143e1a 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -49,10 +49,7 @@ class BlueskyContext: _reference_cache: Dict[Type, Type] def __init__(self, run_engine: Optional[RunEngine] = None) -> None: - if run_engine is None: - run_engine = RunEngine(context_managers=[]) - - self._run_engine = run_engine + self._run_engine = run_engine or RunEngine(context_managers=[]) self._devices = {} self._plans = {} self._plan_functions = {}