diff --git a/src/aare/devices/enum_pv.py b/src/aare/devices/enum_pv.py index ce500b47..4e9ae9f0 100755 --- a/src/aare/devices/enum_pv.py +++ b/src/aare/devices/enum_pv.py @@ -1,126 +1,114 @@ -from typing import cast +import time +from enum import Enum +from typing import Union, Optional +from epics import PV, Motor, poll -from epics import PV - -from .mx_lib import ValueWaitTimeout, pv_wait +class ValueWaitTimeout(Exception): + """Raised when a PV fails to reach a target value within the timeout period.""" + pass -class EnumPv(object): - def __init__(self, args): - mandatory_args = {"pv", "readback", "states"} - supplied_args = set(args.keys()) - if not mandatory_args.issubset(supplied_args): - raise RuntimeError("Missing arguments: %s" % (", ".join(mandatory_args.difference(supplied_args)))) +def pv_wait(pv: Union[PV, Motor], target: Union[str, int, float], timeout: float = 30.0, + polling: float = 0.1, tolerance: float | None = None): + """ + Unified wait function for Motors and PVs (Strings, Enums, Floats). + """ + start_time = time.monotonic() + end_time = start_time + timeout - if type(args["states"]) is not dict: - raise RuntimeError("states must be a dictionary mapping put/readback for waiting purposes") + # Handle Motors + if isinstance(pv, Motor): + if tolerance is None: + # Try to get the motor resolution/deadband + tolerance = pv.get("RDBD") or 0.01 - self.val = PV(args["pv"]) - self.rbv = PV(args["readback"]) + while time.monotonic() < end_time: + if pv.done_moving and abs(pv.readback - target) <= tolerance: + return + poll(polling) - if "name" in args: - self.device_name = args["name"] - else: - self.device_name = self.val.pvname - - if not hasattr(self.val, "enum_strs"): - raise RuntimeError("PV %s does not have enum_strs" % self.val.pvname) - if not hasattr(self.rbv, "enum_strs"): - raise RuntimeError("PV %s does not have enum_strs" % self.rbv.pvname) - - self.__target = None - - p = args.get("move_done_when") - if p is not None: - self.__wait_pv = PV(p[0]) - self.__wait_target = p[1] - else: - self.__wait_pv = None - self.__wait_target = None - - self.states = {} - for k, v in list(args["states"].items()): - self.states[k.upper()] = v.upper() - - self.val_enums = [x.upper() for x in self.val.enum_strs] # type: ignore - self.rbv_enums = [x.upper() for x in self.rbv.enum_strs] # type: ignore - - if "timeout" not in args: - self.__timeout = 60 # seconds - else: - self.__timeout = args["timeout"] - - def __str__(self): - return f"<{self.device_name} at {self.position}>" - - def __repr__(self): - s = ( - f"<{self.device_name} at {self.position} an {self.__class__.__name__} " - f"instance at {hex(id(self))} positions = {self.positions}>" + raise ValueWaitTimeout( + f"Motor {pv.name} timeout. Target: {target}, Current: {pv.readback}, Done: {pv.done_moving}" ) - return s - def has_position(self, position): - return position.upper() in self.states + # Handle PVs (Enums, Strings, Doubles) + pv_type = (pv.type or "").lower() + is_enum = "enum" in pv_type + is_string = "string" in pv_type + + while time.monotonic() < end_time: + # Get current value: as string for Enums/Strings, as raw value for others + as_string = isinstance(target, str) or is_enum or is_string + current = pv.get(as_string=as_string) + + if current is not None: + # Comparison logic + if isinstance(target, str): + if str(current).strip().lower() == target.strip().lower(): + return + elif isinstance(target, (int, float)): + if tolerance is None: + # Default tolerance based on PV precision + prec = getattr(pv, 'precision', 3) or 3 + tolerance = 10 ** -prec + + if abs(float(current) - float(target)) <= tolerance: + return + + poll(polling) + + # If we exit the loop, we timed out + final_val = pv.get(as_string=isinstance(target, str)) + raise ValueWaitTimeout(f"PV {pv.pvname} timeout. Target: {target}, Current: {final_val}") + +def wait_for_movement_to_finish(*motors: Motor, timeout: float = 60.0): + """Wait for a group of motors to stop moving.""" + start_time = time.time() + while not all(m.done_moving for m in motors): + if time.time() - start_time > timeout: + moving = [m.name for m in motors if not m.done_moving] + raise TimeoutError(f"Timeout waiting for motors: {', '.join(moving)}") + poll(0.1) + +class EnumPv: + def __init__(self, pv_name: str, rbv_name: Optional[str] = None, timeout: float = 60.0): + self.control = PV(pv_name) + # If no readback is provided, assume the control PV reflects its own state + self.readback = PV(rbv_name) if rbv_name else self.control + self.timeout = timeout + + # Verify these are actually enums + if not self.control.enum_strs: + raise RuntimeError(f"{pv_name} is not an ENUM PV") @property - def positions(self): - return self.rbv_enums + def position(self) -> str: + """Returns the current string value of the readback.""" + return self.readback.get(as_string=True) - def __position(self) -> str: - return self.rbv.enum_strs[self.rbv.get()] # type: ignore + def get(self) -> int: + """Returns the current integer index of the readback.""" + return int(self.readback.get()) - position = property(__position) + def put(self, value: Union[str, int, Enum], wait: bool = False): + """ + Set the PV value. + Accepts string names, integer indices, or Python Enum members. + """ + # Extract value if a Python Enum was passed + val_to_put = value.value if isinstance(value, Enum) else value - def position_is(self, position) -> bool: - if type(position) in (int, float): - return int(position) == self.get() - else: - return str(position).upper() == str(self.__position()).upper() + self.control.put(val_to_put) - def get(self, req_type=None) -> str | int: - if req_type: - res = self.position - else: - res = cast(int, self.rbv.get()) - return res - - def put(self, value, wait=False) -> None: - self.val.put(value) - self.__target = self.rbv_enums.index(self.states[self.val_enums[value]]) if wait: - self.wait() + self.wait(val_to_put) - def move(self, position, wait=False) -> None: - if type(position) is not int: - if not self.has_position(position): - raise ValueError(f"Invalid position {position} for {self.device_name}") - target = self.val_enums.index(position.upper()) - else: - target = position - - self.put(target, wait=wait) - - def equal_dbr_string(self, val) -> bool: - value = self.rbv.get(as_string=True) - return cast(str, value).lower() == val.lower() - - def equal_dbr_int(self, val) -> bool: - return val == self.rbv.get() - - def wait(self, target=None): - if target is None: - target = self.__target - - if self.__wait_pv: - pv = self.__wait_pv - else: - pv = self.rbv - - if target is None: - target = self.__wait_target + def wait(self, target: Union[str, int, Enum]): + """Blocks until the readback matches the target value.""" + # Convert Enum or name to the format pv_wait expects + wait_val = target.name if isinstance(target, Enum) else target try: - pv_wait(pv, target, timeout=self.__timeout, verbose=True) - except ValueWaitTimeout: - raise ValueWaitTimeout(f"Timeout waiting for device to reach [{target}], currently at [{self.position}]") + pv_wait(self.readback, wait_val, timeout=self.timeout) + except Exception as e: + raise ValueWaitTimeout(f"Timed out waiting for {self.readback.pvname} to reach {wait_val}") from e \ No newline at end of file