diff --git a/.github/actions/install_requirements/action.yml b/.github/actions/install_requirements/action.yml index 25a146d16..aea1e0d4b 100644 --- a/.github/actions/install_requirements/action.yml +++ b/.github/actions/install_requirements/action.yml @@ -30,7 +30,7 @@ runs: - name: Create lockfile run: | mkdir -p lockfiles - pip freeze --exclude-editable > lockfiles/${{ inputs.requirements_file }} + pip freeze --exclude-editable --exclude dodal > lockfiles/${{ inputs.requirements_file }} # delete the self referencing line and make sure it isn't blank sed -i '/file:/d' lockfiles/${{ inputs.requirements_file }} shell: bash @@ -55,4 +55,3 @@ runs: fi fi shell: bash - diff --git a/config/adsim.yaml b/config/adsim.yaml index 83a41f29f..c3b3196f4 100644 --- a/config/adsim.yaml +++ b/config/adsim.yaml @@ -1,2 +1,6 @@ env: - startupScript: blueapi.startup.adsim + sources: + - kind: deviceFunctions + module: blueapi.startup.adsim + - kind: planFunctions + module: blueapi.startup.adsim diff --git a/config/bl45p.yaml b/config/bl45p.yaml index 9463d3e4a..0699c52ea 100644 --- a/config/bl45p.yaml +++ b/config/bl45p.yaml @@ -1,2 +1,6 @@ env: - startupScript: blueapi.startup.bl45p + sources: + - kind: dodal + module: dodal.p45 + - kind: planFunctions + module: blueapi.plans diff --git a/config/defaults.yaml b/config/defaults.yaml index 32d1a3f27..234ac0d70 100644 --- a/config/defaults.yaml +++ b/config/defaults.yaml @@ -1,5 +1,9 @@ env: - startupScript: blueapi.startup.example + sources: + - kind: deviceFunctions + module: blueapi.startup.example + - kind: planFunctions + module: blueapi.plans stomp: host: localhost port: 61613 diff --git a/helm/blueapi/values.yaml b/helm/blueapi/values.yaml index 14110c59d..7f9c8de10 100644 --- a/helm/blueapi/values.yaml +++ b/helm/blueapi/values.yaml @@ -59,7 +59,11 @@ affinity: {} worker: env: - startupScript: blueapi.startup.example + sources: + - kind: deviceFunctions + module: blueapi.startup.example + - kind: planFunctions + module: blueapi.plans stomp: host: activemq port: 61613 diff --git a/pyproject.toml b/pyproject.toml index df52a72d5..3e2976f46 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,8 @@ dependencies = [ "fastapi[all]", "uvicorn", "requests", + "dodal @ git+https://github.com/DiamondLightSource/dodal.git", + ] dynamic = ["version"] license.file = "LICENSE" @@ -70,6 +72,7 @@ write_to = "src/blueapi/_version.py" [tool.mypy] ignore_missing_imports = true # Ignore missing stubs in imported modules +namespace_packages = false # rely only on __init__ files to determine fully qualified module names. [tool.isort] float_to_top = true diff --git a/src/blueapi/config.py b/src/blueapi/config.py index f9243ffbd..8178ea418 100644 --- a/src/blueapi/config.py +++ b/src/blueapi/config.py @@ -1,3 +1,4 @@ +from enum import Enum from pathlib import Path from typing import Any, Generic, Literal, Mapping, Type, TypeVar, Union @@ -9,7 +10,18 @@ LogLevel = Literal["NOTSET", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] -class StompConfig(BlueapiBaseModel): +class SourceKind(str, Enum): + PLAN_FUNCTIONS = "planFunctions" + DEVICE_FUNCTIONS = "deviceFunctions" + DODAL = "dodal" + + +class Source(BaseModel): + kind: SourceKind + module: Union[Path, str] + + +class StompConfig(BaseModel): """ Config for connecting to stomp broker """ @@ -23,11 +35,14 @@ class EnvironmentConfig(BlueapiBaseModel): Config for the RunEngine environment """ - startup_script: Union[Path, str] = "blueapi.startup.example" + sources: list[Source] = [ + Source(kind=SourceKind.DEVICE_FUNCTIONS, module="blueapi.startup.example"), + Source(kind=SourceKind.PLAN_FUNCTIONS, module="blueapi.plans"), + ] def __eq__(self, other: object) -> bool: if isinstance(other, EnvironmentConfig): - return str(self.startup_script) == str(other.startup_script) + return str(self.sources) == str(other.sources) return False diff --git a/src/blueapi/core/context.py b/src/blueapi/core/context.py index 9a1221f7c..61d58aacf 100644 --- a/src/blueapi/core/context.py +++ b/src/blueapi/core/context.py @@ -2,7 +2,6 @@ from dataclasses import dataclass, field from importlib import import_module from inspect import Parameter, signature -from pathlib import Path from types import ModuleType from typing import ( Any, @@ -24,6 +23,7 @@ from pydantic import create_model from pydantic.fields import FieldInfo +from blueapi.config import EnvironmentConfig, SourceKind from blueapi.utils import BlueapiPlanModelConfig, load_module_all from .bluesky_types import ( @@ -73,13 +73,16 @@ def find_device(self, addr: Union[str, List[str]]) -> Optional[Device]: else: return find_component(self.devices, addr) - def with_startup_script(self, path: Union[Path, str]) -> None: - mod = import_module(str(path)) - self.with_module(mod) + def with_config(self, config: EnvironmentConfig) -> None: + for source in config.sources: + mod = import_module(str(source.module)) - def with_module(self, module: ModuleType) -> None: - self.with_plan_module(module) - self.with_device_module(module) + if source.kind is SourceKind.PLAN_FUNCTIONS: + self.with_plan_module(mod) + elif source.kind is SourceKind.DEVICE_FUNCTIONS: + self.with_device_module(mod) + elif source.kind is SourceKind.DODAL: + self.with_dodal_module(mod) def with_plan_module(self, module: ModuleType) -> None: """ @@ -106,9 +109,13 @@ def plan_2(...): self.plan(obj) def with_device_module(self, module: ModuleType) -> None: - for obj in load_module_all(module): - if is_bluesky_compatible_device(obj): - self.device(obj) + self.with_dodal_module(module) + + def with_dodal_module(self, module: ModuleType) -> None: + from dodal.utils import make_all_devices + + for device in make_all_devices(module).values(): + self.device(device) def plan(self, plan: PlanGenerator) -> PlanGenerator: """ diff --git a/src/blueapi/service/handler.py b/src/blueapi/service/handler.py index 173f0551a..c0fd68cd7 100644 --- a/src/blueapi/service/handler.py +++ b/src/blueapi/service/handler.py @@ -22,7 +22,7 @@ def __init__(self, config: Optional[ApplicationConfig] = None) -> None: logging.basicConfig(level=self.config.logging.level) - self.context.with_startup_script(self.config.env.startup_script) + self.context.with_config(self.config.env) self.worker = RunEngineWorker(self.context) self.message_bus = StompMessagingTemplate.autoconfigured(self.config.stomp) diff --git a/src/blueapi/startup/bl45p.py b/src/blueapi/startup/bl45p.py deleted file mode 100644 index 6a8e5e3d7..000000000 --- a/src/blueapi/startup/bl45p.py +++ /dev/null @@ -1,113 +0,0 @@ -from nslsii.ad33 import CamV33Mixin, SingleTriggerV33 -from ophyd import Component as Cpt -from ophyd import DetectorBase, EpicsMotor, MotorBundle -from ophyd.areadetector.base import ADComponent as Cpt -from ophyd.areadetector.cam import AreaDetectorCam -from ophyd.areadetector.detectors import DetectorBase -from ophyd.areadetector.filestore_mixins import FileStoreHDF5, FileStoreIterativeWrite -from ophyd.areadetector.plugins import HDF5Plugin - -from blueapi.plans import * # noqa: F401, F403 - - -class SampleY(MotorBundle): - top: EpicsMotor = Cpt(EpicsMotor, "Y:TOP") - bottom: EpicsMotor = Cpt(EpicsMotor, "Y:BOT") - - -class SampleTheta(MotorBundle): - top: EpicsMotor = Cpt(EpicsMotor, "THETA:TOP") - bottom: EpicsMotor = Cpt(EpicsMotor, "THETA:BOT") - - -class SampleStage(MotorBundle): - x: EpicsMotor = Cpt(EpicsMotor, "X") - y: SampleY = Cpt(SampleY, "") - theta: SampleTheta = Cpt(SampleTheta, "") - - -class Choppers(MotorBundle): - x: EpicsMotor = Cpt(EpicsMotor, "ENDAT") - y: EpicsMotor = Cpt(EpicsMotor, "BISS") - - -_ACQUIRE_BUFFER_PERIOD = 0.2 - - -class NonBlockingCam(AreaDetectorCam, CamV33Mixin): - ... - - -# define a detector device class that has the correct PV suffixes for the rigs -class GigeCamera(SingleTriggerV33, DetectorBase): - class HDF5File(HDF5Plugin, FileStoreHDF5, FileStoreIterativeWrite): - pool_max_buffers = None - file_number_sync = None - file_number_write = None - - def get_frames_per_point(self): - return self.parent.cam.num_images.get() - - cam: NonBlockingCam = Cpt(NonBlockingCam, suffix="DET:") - hdf: HDF5File = Cpt( - HDF5File, - suffix="HDF5:", - root="/dls/tmp/vid18871/data", - write_path_template="%Y", - ) - - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) - self.hdf.kind = "normal" - - # Get stage to wire up the plugins - self.stage_sigs[self.hdf.nd_array_port] = self.cam.port_name.get() - - # Makes the detector allow non-blocking AD plugins but makes Ophyd use - # the AcquireBusy PV to determine when an acquisition is complete - self.cam.ensure_nonblocking() - - # Reset array counter on stage - self.stage_sigs[self.cam.array_counter] = 0 - - # Set image mode to multiple on stage so we have the option, can still - # set num_images to 1 - self.stage_sigs[self.cam.image_mode] = "Multiple" - - # For now, this Ophyd device does not support hardware - # triggered scanning, disable on stage - self.stage_sigs[self.cam.trigger_mode] = "Off" - - def make_data_key(self): - source = "PV:{}".format(self.prefix) - # This shape is expected to match arr.shape for the array. - shape = ( - self.cam.num_images.get(), - self.cam.array_size.array_size_y.get(), - self.cam.array_size.array_size_x.get(), - ) - return dict(shape=shape, source=source, dtype="array", external="FILESTORE:") - - def stage(self, *args, **kwargs): - # We have to manually set the acquire period bcause the EPICS driver - # doesn't do it for us. If acquire time is a staged signal, we use the - # stage value to calculate the acquire period, otherwise we perform - # a caget and use the current acquire time. - if self.cam.acquire_time in self.stage_sigs: - acquire_time = self.stage_sigs[self.cam.acquire_time] - else: - acquire_time = self.cam.acquire_time.get() - acquire_period = acquire_time + _ACQUIRE_BUFFER_PERIOD - self.stage_sigs[self.cam.acquire_period] = acquire_period - - # Now calling the super method should set the acquire period - super(GigeCamera, self).stage(*args, **kwargs) - - -sample = SampleStage(name="sample", prefix="BL45P-MO-STAGE-01:") -choppers = Choppers(name="chopper", prefix="BL45P-MO-CHOP-01:") -det = GigeCamera(name="det", prefix="BL45P-EA-MAP-01:") -diff = GigeCamera(name="diff", prefix="BL45P-EA-DIFF-01:") - -for device in sample, choppers, det, diff: - device.wait_for_connection() # type: ignore diff --git a/src/blueapi/startup/example.py b/src/blueapi/startup/example.py index fbc27768e..2244071d7 100644 --- a/src/blueapi/startup/example.py +++ b/src/blueapi/startup/example.py @@ -1,38 +1,74 @@ from ophyd.sim import Syn2DGauss, SynGauss, SynSignal -from blueapi.plans import * # noqa: F401, F403 - from .simmotor import BrokenSynAxis, SynAxisWithMotionEvents -x = SynAxisWithMotionEvents(name="x", delay=1.0, events_per_move=8) -y = SynAxisWithMotionEvents(name="y", delay=3.0, events_per_move=24) -z = SynAxisWithMotionEvents(name="z", delay=2.0, events_per_move=16) -theta = SynAxisWithMotionEvents( - name="theta", delay=0.2, events_per_move=12, egu="degrees" -) -x_err = BrokenSynAxis(name="x_err", timeout=1.0) -sample_pressure = SynAxisWithMotionEvents( - name="sample_pressure", delay=30.0, events_per_move=128, egu="MPa", value=0.101 -) -sample_temperature = SynSignal( - func=lambda: ((x.position + y.position + z.position) / 1000.0) + 20.0, + +def x(name="x") -> SynAxisWithMotionEvents: + return SynAxisWithMotionEvents(name=name, delay=1.0, events_per_move=8) + + +def y(name="y") -> SynAxisWithMotionEvents: + return SynAxisWithMotionEvents(name=name, delay=3.0, events_per_move=24) + + +def z(name="z") -> SynAxisWithMotionEvents: + return SynAxisWithMotionEvents(name=name, delay=2.0, events_per_move=16) + + +def theta(name="theta") -> SynAxisWithMotionEvents: + return SynAxisWithMotionEvents( + name=name, delay=0.2, events_per_move=12, egu="degrees" + ) + + +def x_err(name="x_err") -> BrokenSynAxis: + return BrokenSynAxis(name=name, timeout=1.0) + + +def sample_pressure(name="sample_pressure") -> SynAxisWithMotionEvents: + return SynAxisWithMotionEvents( + name=name, delay=30.0, events_per_move=128, egu="MPa", value=0.101 + ) + + +def sample_temperature( + x: SynAxisWithMotionEvents, + y: SynAxisWithMotionEvents, + z: SynAxisWithMotionEvents, name="sample_temperature", -) -image_det = Syn2DGauss( +) -> SynSignal: + return SynSignal( + func=lambda: ((x.position + y.position + z.position) / 1000.0) + 20.0, + name=name, + ) + + +def image_det( + x: SynAxisWithMotionEvents, + y: SynAxisWithMotionEvents, name="image_det", - motor0=x, - motor_field0="x", - motor1=y, - motor_field1="y", - center=(0, 0), - Imax=1, - labels={"detectors"}, -) -current_det = SynGauss( +) -> Syn2DGauss: + return Syn2DGauss( + name=name, + motor0=x, + motor_field0="x", + motor1=y, + motor_field1="y", + center=(0, 0), + Imax=1, + labels={"detectors"}, + ) + + +def current_det( + x: SynAxisWithMotionEvents, name="current_det", - motor=x, - motor_field="x", - center=0.0, - Imax=1, - labels={"detectors"}, -) +) -> SynGauss: + return SynGauss( + name=name, + motor=x, + motor_field="x", + center=0.0, + Imax=1, + labels={"detectors"}, + ) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/core/fake_device_module.py b/tests/core/fake_device_module.py new file mode 100644 index 000000000..feaafac55 --- /dev/null +++ b/tests/core/fake_device_module.py @@ -0,0 +1,31 @@ +from unittest.mock import MagicMock + +from ophyd import EpicsMotor + + +def fake_motor_bundle_b( + fake_motor_x: EpicsMotor, + fake_motor_y: EpicsMotor, +) -> EpicsMotor: + return _mock_with_name("motor_bundle_b") + + +def fake_motor_x() -> EpicsMotor: + return _mock_with_name("motor_x") + + +def fake_motor_y() -> EpicsMotor: + return _mock_with_name("motor_y") + + +def fake_motor_bundle_a( + fake_motor_x: EpicsMotor, + fake_motor_y: EpicsMotor, +) -> EpicsMotor: + return _mock_with_name("motor_bundle_a") + + +def _mock_with_name(name: str) -> MagicMock: + mock = MagicMock() + mock.name = name + return mock diff --git a/tests/core/fake_plan_module.py b/tests/core/fake_plan_module.py new file mode 100644 index 000000000..450bc5831 --- /dev/null +++ b/tests/core/fake_plan_module.py @@ -0,0 +1 @@ +from blueapi.plans import scan # noqa: F401 diff --git a/tests/core/test_context.py b/tests/core/test_context.py index cca1145e5..b5cc0436d 100644 --- a/tests/core/test_context.py +++ b/tests/core/test_context.py @@ -6,6 +6,7 @@ from bluesky.protocols import Descriptor, Movable, Readable, Reading, SyncOrAsync from ophyd.sim import SynAxis, SynGauss +from blueapi.config import EnvironmentConfig, Source, SourceKind from blueapi.core import ( BlueskyContext, MsgGenerator, @@ -110,6 +111,13 @@ def test_add_invalid_plan(empty_context: BlueskyContext, plan: PlanGenerator) -> empty_context.plan(plan) +def test_add_plan_from_module(empty_context: BlueskyContext) -> None: + import tests.core.fake_plan_module as plan_module + + empty_context.with_plan_module(plan_module) + assert {"scan"} == empty_context.plans.keys() + + 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 @@ -137,6 +145,18 @@ def test_override_device_name( assert empty_context.devices["foo"] is sim_motor +def test_add_devices_from_module(empty_context: BlueskyContext) -> None: + import tests.core.fake_device_module as device_module + + empty_context.with_device_module(device_module) + assert { + "motor_x", + "motor_y", + "motor_bundle_a", + "motor_bundle_b", + } == empty_context.devices.keys() + + @pytest.mark.parametrize( "addr", ["sim", "sim_det", "sim.setpoint", ["sim"], ["sim", "setpoint"]] ) @@ -170,6 +190,31 @@ def test_add_non_device(empty_context: BlueskyContext) -> None: empty_context.device("not a device") # type: ignore +def test_add_devices_and_plans_from_modules_with_config( + empty_context: BlueskyContext, +) -> None: + empty_context.with_config( + EnvironmentConfig( + sources=[ + Source( + kind=SourceKind.DEVICE_FUNCTIONS, + module="tests.core.fake_device_module", + ), + Source( + kind=SourceKind.PLAN_FUNCTIONS, module="tests.core.fake_plan_module" + ), + ] + ) + ) + assert { + "motor_x", + "motor_y", + "motor_bundle_a", + "motor_bundle_b", + } == empty_context.devices.keys() + assert {"scan"} == empty_context.plans.keys() + + def test_function_spec(empty_context: BlueskyContext) -> None: spec = empty_context._type_spec_for_function(has_some_params) assert spec["foo"][0] == int diff --git a/tests/utils/test_modules.py b/tests/utils/test_modules.py index 429f035e4..2a33cc17a 100644 --- a/tests/utils/test_modules.py +++ b/tests/utils/test_modules.py @@ -4,10 +4,10 @@ def test_imports_all(): - module = import_module(".hasall", package="utils") + module = import_module(".hasall", package="tests.utils") assert list(load_module_all(module)) == ["hello", 9] def test_imports_everything_without_all(): - module = import_module(".lacksall", package="utils") + module = import_module(".lacksall", package="tests.utils") assert list(load_module_all(module)) == [3, "hello", 9] diff --git a/tests/worker/test_reworker.py b/tests/worker/test_reworker.py index 3f53b2a0e..458b0fe88 100644 --- a/tests/worker/test_reworker.py +++ b/tests/worker/test_reworker.py @@ -4,6 +4,7 @@ import pytest +from blueapi.config import EnvironmentConfig from blueapi.core import BlueskyContext, EventStream from blueapi.worker import ( RunEngineWorker, @@ -20,7 +21,7 @@ @pytest.fixture def context() -> BlueskyContext: ctx = BlueskyContext() - ctx.with_startup_script("blueapi.startup.example") + ctx.with_config(EnvironmentConfig()) return ctx