diff --git a/xtreme_bec/devices/uind_mode.py b/xtreme_bec/devices/uind_mode.py new file mode 100644 index 0000000..75f1e5f --- /dev/null +++ b/xtreme_bec/devices/uind_mode.py @@ -0,0 +1,71 @@ +from enum import Enum +from typing import Literal + +from ophyd import Component as Cpt +from ophyd import Device, EpicsSignal, Kind, StatusBase + + +class UIndModeEnum(str, Enum): + OFF = "Off" + LINEAR_H = "Linear H" + LINEAR_V_PLUS = "Linear V+" + LINEAR_V_MINUS = "Linear V-" + CIRCULAR_PLUS = "Circular+" + CIRCULAR_MINUS = "Circular-" + PLUS_45_DEG = "+45 deg" + MINUS_45_DEG = "-45 deg" + + +class UIndMode(Device): + USER_ACCESS = ["set_mode"] + + value = Cpt(EpicsSignal, "X07MA-UIND:MODE", auto_monitor=True, kind=Kind.hinted) + value_string = Cpt(EpicsSignal, "X07MA-UIND:MODE", auto_monitor=True, string=True) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.value.name = self.name + + def set_mode( + self, + mode: ( + Literal[ + "Off", + "Linear H", + "Linear V+", + "Linear V-", + "Circular+", + "Circular-", + "+45 deg", + "-45 deg", + ] + | int + ), + ) -> StatusBase: + """ + Set the undulator polarization mode. + + Args: + mode (Literal["Off", "Linear H", "Linear V+", "Linear V-", "Circular+", "Circular-", "+45 deg", "-45 deg"] | int): The desired polarization mode, either as a string or an integer index. + """ + if isinstance(mode, str): + try: + mode_enum = UIndModeEnum(mode) + mode_val = mode_enum.value + except ValueError: + # pylint: disable=raise-missing-from + raise ValueError( + f"Invalid mode string: {mode}. Valid options are: {[m.value for m in UIndModeEnum]}" + ) + elif isinstance(mode, int): + try: + mode_enum = list(UIndModeEnum)[mode] + mode_val = mode_enum.value + except IndexError: + # pylint: disable=raise-missing-from + raise ValueError( + f"Invalid mode index: {mode}. Valid indices are: 0 to {len(UIndModeEnum) - 1}, corresponding to: {[m.value for m in UIndModeEnum]}" + ) + else: + raise TypeError("Mode must be either a string or an integer index.") + return self.value.set(mode_val)