11import abc
22import asyncio
33from dataclasses import dataclass
4- from typing import Any
4+ from typing import Any , Generic , TypeVar
55
66import numpy as np
77from bluesky .protocols import Movable
2121
2222from dodal .log import LOGGER
2323
24+ T = TypeVar ("T" )
25+
2426
2527class UndulatorGateStatus (StrictEnum ):
2628 OPEN = "Open"
@@ -99,7 +101,7 @@ async def estimate_motor_timeout(
99101 return abs ((target_pos - cur_pos ) * 2.0 / vel ) + 1
100102
101103
102- class SafeUndulatorMover (StandardReadable , Movable ):
104+ class SafeUndulatorMover (StandardReadable , Movable [ T ], Generic [ T ] ):
103105 """A device that will check it's safe to move the undulator before moving it and
104106 wait for the undulator to be safe again before calling the move complete.
105107 """
@@ -115,7 +117,7 @@ def __init__(self, set_move: SignalW, prefix: str, name: str = ""):
115117 super ().__init__ (name )
116118
117119 @AsyncStatus .wrap
118- async def set (self , value ) -> None :
120+ async def set (self , value : T ) -> None :
119121 LOGGER .info (f"Setting { self .name } to { value } " )
120122 await self .raise_if_cannot_move ()
121123 await self ._set_demand_positions (value )
@@ -125,7 +127,7 @@ async def set(self, value) -> None:
125127 await wait_for_value (self .gate , UndulatorGateStatus .CLOSE , timeout = timeout )
126128
127129 @abc .abstractmethod
128- async def _set_demand_positions (self , value ) -> None :
130+ async def _set_demand_positions (self , value : T ) -> None :
129131 """Set the demand positions on the device without actually hitting move."""
130132
131133 @abc .abstractmethod
@@ -139,7 +141,7 @@ async def raise_if_cannot_move(self) -> None:
139141 raise RuntimeError (f"{ self .name } is already in motion." )
140142
141143
142- class UndulatorGap (SafeUndulatorMover ):
144+ class UndulatorGap (SafeUndulatorMover [ float ] ):
143145 """A device with a collection of epics signals to set Apple 2 undulator gap motion.
144146 Only PV used by beamline are added the full list is here:
145147 /dls_sw/work/R3.14.12.7/support/insertionDevice/db/IDGapVelocityControl.template
@@ -185,7 +187,7 @@ def __init__(self, prefix: str, name: str = ""):
185187 self .user_readback = epics_signal_r (float , prefix + "CURRGAPD" )
186188 super ().__init__ (self .set_move , prefix , name )
187189
188- async def _set_demand_positions (self , value ) -> None :
190+ async def _set_demand_positions (self , value : float ) -> None :
189191 await self .user_setpoint .set (str (value ))
190192
191193 async def get_timeout (self ) -> float :
@@ -234,7 +236,7 @@ def __init__(self, prefix: str, infix: str, name: str = ""):
234236 super ().__init__ (name = name )
235237
236238
237- class UndulatorPhaseAxes (SafeUndulatorMover ):
239+ class UndulatorPhaseAxes (SafeUndulatorMover [ Apple2PhasesVal ] ):
238240 """
239241 A collection of 4 phase Motor to make up the full id phase motion. We are using the diamond pv convention.
240242 e.g. top_outer == Q1
@@ -290,7 +292,7 @@ async def get_timeout(self) -> float:
290292 return np .max (timeouts )
291293
292294
293- class UndulatorJawPhase (SafeUndulatorMover ):
295+ class UndulatorJawPhase (SafeUndulatorMover [ float ] ):
294296 """
295297 A JawPhase movable, this is use for moving the jaw phase which is use to control the
296298 linear arbitrary polarisation but only one some of the beamline.
0 commit comments