diff --git a/src/dodal/devices/electron_analyser/energy_sources.py b/src/dodal/devices/electron_analyser/energy_sources.py index 4b6affeb72..2a09ea15aa 100644 --- a/src/dodal/devices/electron_analyser/energy_sources.py +++ b/src/dodal/devices/electron_analyser/energy_sources.py @@ -1,5 +1,3 @@ -from abc import abstractmethod - from ophyd_async.core import ( Reference, SignalR, @@ -11,26 +9,10 @@ ) from dodal.devices.electron_analyser.enums import SelectedSource +from dodal.protocols import EnergyWrapper -class AbstractEnergySource(StandardReadable): - """ - Abstract device that wraps an energy source signal and provides common interface via - a energy signal. - """ - - def __init__(self, name: str = "") -> None: - super().__init__(name) - - @property - @abstractmethod - def energy(self) -> SignalR[float]: - """ - Signal to provide the excitation energy value in eV. - """ - - -class EnergySource(AbstractEnergySource): +class EnergySource(StandardReadable, EnergyWrapper[SignalR[float]]): """ Wraps a signal that relates to energy and provides common interface via energy signal. It provides the name of the wrapped signal as a child signal in the @@ -51,7 +33,7 @@ def energy(self) -> SignalR[float]: return self._source_ref() -class DualEnergySource(AbstractEnergySource): +class DualEnergySource(StandardReadable, EnergyWrapper[SignalR[float]]): """ Holds two EnergySource devices and provides a signal to read energy depending on which source is selected. This is controlled by a selected_source signal which can diff --git a/src/dodal/devices/pgm.py b/src/dodal/devices/pgm.py index c0abadf2a6..5750227617 100644 --- a/src/dodal/devices/pgm.py +++ b/src/dodal/devices/pgm.py @@ -6,8 +6,10 @@ from ophyd_async.epics.core import epics_signal_rw from ophyd_async.epics.motor import Motor +from dodal.protocols import EnergySource -class PGM(StandardReadable): + +class PGM(StandardReadable, EnergySource[Motor]): """ Plane grating monochromator, it is use in soft x-ray beamline to generate monochromic beam. """ diff --git a/src/dodal/protocols.py b/src/dodal/protocols.py new file mode 100644 index 0000000000..4d86697c70 --- /dev/null +++ b/src/dodal/protocols.py @@ -0,0 +1,28 @@ +from abc import abstractmethod +from typing import Generic, Protocol, TypeVar + +from ophyd_async.core import Device + +DeviceT = TypeVar("DeviceT", bound=Device) + + +class EnergySource(Protocol, Generic[DeviceT]): + """Class that directly defines an energy device.""" + + energy: DeviceT + + +DeviceT_CO = TypeVar("DeviceT_CO", bound=Device, covariant=True) + + +class EnergyWrapper(Protocol, Generic[DeviceT_CO]): + """Class that indirectly defines an energy device and exposes it via a property.""" + + @property + @abstractmethod + def energy(self) -> DeviceT_CO: + """Property or class variable that defines energy in some way.""" + + +"""A class that defines an energy device.""" +EnergyDevice = EnergyWrapper[DeviceT] | EnergySource[DeviceT]