Create energy protocol

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

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.
     """

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]