diff --git a/ophyd_devices/interfaces/base_classes/psi_device_base.py b/ophyd_devices/interfaces/base_classes/psi_device_base.py index 9d0bdf7..c5b19e8 100644 --- a/ophyd_devices/interfaces/base_classes/psi_device_base.py +++ b/ophyd_devices/interfaces/base_classes/psi_device_base.py @@ -8,10 +8,11 @@ import inspect import time from typing import TYPE_CHECKING, Callable -from ophyd import Device, DeviceStatus, Staged, StatusBase +from ophyd import Device, Staged +from ophyd.status import StatusBase as OphydStatusBase from ophyd_devices.tests.utils import get_mock_scan_info -from ophyd_devices.utils.psi_device_base_utils import FileHandler, TaskHandler +from ophyd_devices.utils.psi_device_base_utils import DeviceStatus, FileHandler, TaskHandler if TYPE_CHECKING: # pragma: no cover from bec_lib.devicemanager import DeviceManagerBase, ScanInfo @@ -53,7 +54,9 @@ class PSIDeviceBase(Device): Args: name (str) : Name of the device + prefix (str): The prefix for the device. scan_info (ScanInfo): The scan info to use. + device_manager (DeviceManagerBase): The device manager to use. """ # Make sure device_manager is not passed to super().__init__ if not specified # This is to avoid issues with ophyd.OphydObject.__init__ when the parent is ophyd.Device @@ -64,8 +67,15 @@ class PSIDeviceBase(Device): super().__init__(device_manager=device_manager, prefix=prefix, name=name, **kwargs) else: super().__init__(prefix=prefix, name=name, **kwargs) + timeout = kwargs.pop("timeout", None) + if getattr(self, "_timeout", None) is None: + self._timeout = self._normalize_timeout(timeout) + if not isinstance(self._timeout, (float, int, type(None))): + raise TypeError(f"Timeout must be a float, int or None, got {type(self._timeout)}") + self._timeout = self._normalize_timeout(timeout) + self._set_timeout_signal(timeout) self._stopped = False - self._stoppable_status_objects: list[StatusBase] = [] + self._stoppable_status_objects: list[OphydStatusBase] = [] self.task_handler = TaskHandler(parent=self) self.file_utils = FileHandler() if scan_info is None: @@ -96,54 +106,69 @@ class PSIDeviceBase(Device): def stopped(self, value: bool): self._stopped = value + @staticmethod + def _normalize_timeout(timeout: float | int | None) -> float | None: + """Normalize non-positive timeout values to no timeout.""" + if not isinstance(timeout, (float, int, type(None))): + raise TypeError(f"Timeout must be a float, int, or None, got {type(timeout)}") + if timeout is None or timeout <= 0: + return None + return float(timeout) + + def _set_timeout_signal(self, timeout: float | None) -> None: + """Initialize an optional timeout component from the constructor value.""" + if timeout is None or "timeout" not in getattr(self, "component_names", ()): + return + self.timeout.put(0 if timeout <= 0 else timeout) + ######################################## # Wrapper around Device class methods # ######################################## - def stage(self) -> list[object] | DeviceStatus | StatusBase: # type: ignore + def stage(self) -> list[object] | DeviceStatus | OphydStatusBase: # type: ignore """Stage the device.""" if self.staged != Staged.no: return super().stage() self.stopped = False super_staged = super().stage() status = self.on_stage() # pylint: disable=assignment-from-no-return - if isinstance(status, StatusBase): + if isinstance(status, OphydStatusBase): return status return super_staged - def unstage(self) -> list[object] | DeviceStatus | StatusBase: # type: ignore + def unstage(self) -> list[object] | DeviceStatus | OphydStatusBase: # type: ignore """Unstage the device.""" super_unstage = super().unstage() status = self.on_unstage() # pylint: disable=assignment-from-no-return self._stop_stoppable_status_objects() - if isinstance(status, StatusBase): + if isinstance(status, OphydStatusBase): return status return super_unstage - def pre_scan(self) -> DeviceStatus | StatusBase | None: + def pre_scan(self) -> DeviceStatus | OphydStatusBase | None: """Pre-scan function.""" status = self.on_pre_scan() # pylint: disable=assignment-from-no-return return status - def trigger(self) -> DeviceStatus | StatusBase: + def trigger(self) -> DeviceStatus | OphydStatusBase: """Trigger the device.""" super_trigger = super().trigger() status = self.on_trigger() # pylint: disable=assignment-from-no-return return status if status else super_trigger - def complete(self) -> DeviceStatus | StatusBase: + def complete(self) -> DeviceStatus | OphydStatusBase: """Complete the device.""" status = self.on_complete() # pylint: disable=assignment-from-no-return - if isinstance(status, StatusBase): + if isinstance(status, OphydStatusBase): return status status = DeviceStatus(self) status.set_finished() return status - def kickoff(self) -> DeviceStatus | StatusBase: + def kickoff(self) -> DeviceStatus | OphydStatusBase: """Kickoff the device.""" status = self.on_kickoff() # pylint: disable=assignment-from-no-return - if isinstance(status, StatusBase): + if isinstance(status, OphydStatusBase): return status status = DeviceStatus(self) status.set_finished() @@ -172,15 +197,15 @@ class PSIDeviceBase(Device): # Stoppable Status Objects Management # ######################################## - def cancel_on_stop(self, status: StatusBase) -> None: + def cancel_on_stop(self, status: OphydStatusBase) -> None: """ Register a status object to be cancelled when the device is stopped. Args: - status (StatusBase): The status object to be cancelled. + status (OphydStatusBase): The status object to be cancelled. """ - if not isinstance(status, StatusBase): - raise TypeError("status must be an instance of StatusBase") + if not isinstance(status, OphydStatusBase): + raise TypeError("status must be an instance of ophyd.StatusBase") self._stoppable_status_objects.append(status) def _clear_stoppable_status_objects(self) -> None: @@ -258,26 +283,26 @@ class PSIDeviceBase(Device): Default values for signals should be set here. """ - def on_stage(self) -> DeviceStatus | StatusBase | None: + def on_stage(self) -> DeviceStatus | OphydStatusBase | None: """ Called while staging the device. Information about the upcoming scan can be accessed from the scan_info (self.scan_info.msg) object. """ - def on_unstage(self) -> DeviceStatus | StatusBase | None: + def on_unstage(self) -> DeviceStatus | OphydStatusBase | None: """Called while unstaging the device.""" - def on_pre_scan(self) -> DeviceStatus | StatusBase | None: + def on_pre_scan(self) -> DeviceStatus | OphydStatusBase | None: """Called right before the scan starts on all devices automatically.""" - def on_trigger(self) -> DeviceStatus | StatusBase | None: + def on_trigger(self) -> DeviceStatus | OphydStatusBase | None: """Called when the device is triggered.""" - def on_complete(self) -> DeviceStatus | StatusBase | None: + def on_complete(self) -> DeviceStatus | OphydStatusBase | None: """Called to inquire if a device has completed a scans.""" - def on_kickoff(self) -> DeviceStatus | StatusBase | None: + def on_kickoff(self) -> DeviceStatus | OphydStatusBase | None: """Called to kickoff a device for a fly scan. Has to be called explicitly.""" def on_stop(self) -> None: diff --git a/ophyd_devices/utils/psi_device_base_utils.py b/ophyd_devices/utils/psi_device_base_utils.py index 0cd20bf..7e37e7d 100644 --- a/ophyd_devices/utils/psi_device_base_utils.py +++ b/ophyd_devices/utils/psi_device_base_utils.py @@ -204,6 +204,32 @@ def _run_callbacks_with_diagnostics( status._callbacks.clear() +def _get_default_timeout(obj) -> float | None: + """ + Walk an object/parent chain to find a PSIDeviceBase default timeout. + We have to filter down to instances of PSIDeviceBase to avoid conflicting + with timeout attributes set on other objects in the device hierarchy, such as + regular ophyd signals. + """ + from ophyd_devices.interfaces.base_classes.psi_device_base import PSIDeviceBase + + current = obj + seen: set[int] = set() + while current is not None and id(current) not in seen: + seen.add(id(current)) + if isinstance(current, PSIDeviceBase): + return getattr(current, "_timeout") + current = getattr(current, "parent", None) + return None + + +def _resolve_timeout(timeout, obj) -> float | None: + """Use an explicit timeout, falling back to the object's default timeout.""" + if timeout is not None: + return timeout + return _get_default_timeout(obj) + + class StatusBase(_StatusBase): """Base class for all status objects.""" @@ -221,6 +247,7 @@ class StatusBase(_StatusBase): ): self.obj = obj self._timeout_diagnostics = _StatusTimeoutDiagnostics(description=description) + timeout = _resolve_timeout(timeout, obj) super().__init__(timeout=timeout, settle_time=settle_time, done=done, success=success) self._timeout_diagnostics.bind(self) @@ -354,6 +381,7 @@ class Status(_Status): description: str | None = None, ): self._timeout_diagnostics = _StatusTimeoutDiagnostics(description=description) + timeout = _resolve_timeout(timeout, obj) super().__init__( obj=obj, timeout=timeout, settle_time=settle_time, done=done, success=success ) @@ -370,9 +398,21 @@ class Status(_Status): class DeviceStatus(_DeviceStatus): """Thin wrapper around DeviceStatus to add __and__ operator.""" - def __init__(self, device, description: str | None = None, **kwargs): + def __init__( + self, + device, + *, + timeout=None, + settle_time=0, + done=None, + success=None, + description: str | None = None, + ): self._timeout_diagnostics = _StatusTimeoutDiagnostics(description=description) - super().__init__(device=device, **kwargs) + timeout = _resolve_timeout(timeout, device) + super().__init__( + device=device, timeout=timeout, settle_time=settle_time, done=done, success=success + ) self._timeout_diagnostics.bind(self) def __and__(self, other): @@ -387,10 +427,28 @@ class MoveStatus(_MoveStatus): """Thin wrapper around MoveStatus to ensure __and__ operator and stop on failure.""" def __init__( - self, positioner, target, *, start_ts=None, description: str | None = None, **kwargs + self, + positioner, + target, + *, + start_ts=None, + timeout=None, + settle_time=0, + done=None, + success=None, + description: str | None = None, ): self._timeout_diagnostics = _StatusTimeoutDiagnostics(description=description) - super().__init__(positioner=positioner, target=target, start_ts=start_ts, **kwargs) + timeout = _resolve_timeout(timeout, positioner) + super().__init__( + positioner=positioner, + target=target, + start_ts=start_ts, + timeout=timeout, + settle_time=settle_time, + done=done, + success=success, + ) self._timeout_diagnostics.bind(self) def __and__(self, other): diff --git a/tests/test_psi_device_base.py b/tests/test_psi_device_base.py index ab1c388..24bcd12 100644 --- a/tests/test_psi_device_base.py +++ b/tests/test_psi_device_base.py @@ -5,7 +5,8 @@ import time from unittest import mock import pytest -from ophyd import Device +from ophyd import Component as Cpt +from ophyd import Device, Signal from ophyd.status import StatusBase from ophyd_devices.interfaces.base_classes.psi_device_base import DeviceStoppedError, PSIDeviceBase @@ -24,6 +25,19 @@ class SimDevice(PSIDeviceBase, Device): """Simulated Device with PSI Device Base""" +class TimeoutSignalDevice(PSIDeviceBase, Device): + """Device that exposes the base timeout as a signal.""" + + timeout = Cpt(Signal, value=10) + + def __init__(self, timeout=10, **kwargs): + super().__init__(timeout=timeout, **kwargs) + self.timeout.subscribe(self._on_timeout_change, run=False) + + def _on_timeout_change(self, value, **kwargs): + self._timeout = self._normalize_timeout(value) + + @pytest.fixture def device_positioner(): """Fixture for Device""" @@ -82,6 +96,36 @@ def test_psi_device_base_init_with_device_manager(): assert device_2.device_manager is dm +def test_psi_device_base_timeout_init_arg(): + """Test default timeout initialization.""" + assert SimDevice(name="device")._timeout is None + assert SimDevice(name="device", timeout=3)._timeout == 3 + assert SimDevice(name="device", timeout=0)._timeout is None + + +def test_psi_device_base_timeout_signal_compatibility(): + """Test that subclasses can expose timeout as a signal.""" + device = TimeoutSignalDevice(name="device") + + assert device._timeout == 10 + assert device.timeout.get() == 10 + + device.timeout.set(5).wait() + assert device._timeout == 5 + + device.timeout.set(0).wait() + assert device.timeout.get() == 0 + assert device._timeout is None + + +def test_psi_device_base_fallback_statuses_use_default_timeout(): + """Test fallback complete and kickoff statuses use the base timeout.""" + device = SimDevice(name="device", timeout=3) + + assert device.complete().timeout == 3 + assert device.kickoff().timeout == 3 + + def test_on_stage_hook(device): """Test user method hooks""" with mock.patch.object(device, "on_stage") as mock_on_stage: diff --git a/tests/test_psi_motors.py b/tests/test_psi_motors.py index 28e2d2a..cf7ef75 100644 --- a/tests/test_psi_motors.py +++ b/tests/test_psi_motors.py @@ -179,6 +179,15 @@ def test_epics_motor_move_timeout_includes_initialization_traceback(mock_epics_m assert isinstance(exc_info.value, ExceptionWithErrorInfo) +def test_epics_user_motor_vme_timeout_init_arg_updates_signal(): + """PSIDeviceBase timeout init arg should initialize the EpicsMotor timeout signal.""" + with patched_device( + EpicsUserMotorVME, _mock_pv_initial_value=2, name="motor", timeout=4 + ) as motor: + assert motor._timeout == 4 + assert motor.timeout.get() == 4 + + @pytest.fixture(scope="function") def motor(): with patched_device(EpicsUserMotorVME, _mock_pv_initial_value=2, name="motor") as mtr: diff --git a/tests/test_psi_positioner.py b/tests/test_psi_positioner.py index 7fd2a80..5ac309f 100644 --- a/tests/test_psi_positioner.py +++ b/tests/test_psi_positioner.py @@ -104,6 +104,24 @@ def test_status_completed_when_req_done_sub_runs(mock_psi_positioner: PSISimpleP assert st.done +def test_positioner_timeout_init_arg_survives_positioner_base_init(): + """PSIDeviceBase timeout should remain the default move timeout.""" + pos = PSISimplePositioner(name="positioner", prefix="SIM:MOTOR", deadband=0.0013, timeout=3) + + assert pos._timeout == 3 + + +def test_positioner_move_uses_timeout_init_arg(mock_psi_positioner: PSISimplePositioner): + """Positioner move statuses should inherit the init timeout by default.""" + mock_psi_positioner._timeout = 3 + mock_psi_positioner.motor_done_move._read_pv.mock_data = 0 + mock_psi_positioner._position = 0 + + st = mock_psi_positioner.move(1, wait=False) + + assert st.timeout == 3 + + def test_mdm_used_for_moving_if_available(mock_psi_positioner): mock_psi_positioner.wait_for_connection() mock_psi_positioner.motor_done_move._read_pv.mock_data = 0 diff --git a/tests/test_utils.py b/tests/test_utils.py index a25bb35..e66a7af 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -18,6 +18,7 @@ from ophyd.status import WaitTimeoutError from typeguard import TypeCheckError from ophyd_devices.devices.psi_motor import EpicsMotor +from ophyd_devices.interfaces.base_classes.psi_device_base import PSIDeviceBase from ophyd_devices.tests.utils import MockPV, patched_device from ophyd_devices.utils.bec_signals import ( AsyncMultiSignal, @@ -1150,6 +1151,72 @@ class Positioner(Device): pass +class TimeoutDevice(PSIDeviceBase, Device): + sig = Cpt(Signal, value=0) + + def __init__(self, *, timeout=3, **kwargs): + super().__init__(timeout=timeout, **kwargs) + + +class TimeoutPositioner(PSIDeviceBase, Positioner): + def __init__(self, *, timeout=3, **kwargs): + super().__init__(timeout=timeout, **kwargs) + + +class PlainTimeoutDevice(Device): + def __init__(self, *, timeout=3, **kwargs): + self._timeout = timeout + super().__init__(**kwargs) + + +def test_patched_status_types_use_default_timeout_from_object(): + """Patched status objects should use the object's default timeout.""" + dev = TimeoutDevice(name="device", timeout=3) + pos = TimeoutPositioner(name="positioner", timeout=4) + + statuses = [ + StatusBase(obj=dev), + Status(obj=dev), + DeviceStatus(dev), + MoveStatus(pos, target=10), + TaskStatus(obj=dev), + SubscriptionStatus(dev.sig, callback=lambda *args, **kwargs: False, run=False), + CompareStatus(dev.sig, value=1, run=False), + ExceptionStatus(dev.sig, value=1, run=False), + TransitionStatus(dev.sig, transitions=[1], run=False), + ] + + assert [status.timeout for status in statuses] == [3, 3, 3, 4, 3, 3, 3, 3, 3] + + +def test_patched_status_explicit_timeout_overrides_object_default(): + """Explicit status timeouts should take precedence over object defaults.""" + dev = TimeoutDevice(name="device", timeout=3) + pos = TimeoutPositioner(name="positioner", timeout=4) + + assert StatusBase(obj=dev, timeout=7).timeout == 7 + assert Status(obj=dev, timeout=7).timeout == 7 + assert DeviceStatus(dev, timeout=7).timeout == 7 + assert MoveStatus(pos, target=10, timeout=7).timeout == 7 + assert CompareStatus(dev.sig, value=1, timeout=7, run=False).timeout == 7 + + +def test_patched_status_ignores_signal_internal_timeout(): + """Signal connection timeouts should not become status completion timeouts.""" + sig = Signal(name="signal", value=0) + sig._timeout = 2 + + assert CompareStatus(sig, value=1, run=False).timeout is None + + +def test_patched_status_ignores_plain_device_internal_timeout(): + """Only PSIDeviceBase instances should provide default status timeouts.""" + dev = PlainTimeoutDevice(name="device", timeout=3) + + assert StatusBase(obj=dev).timeout is None + assert DeviceStatus(dev).timeout is None + + @pytest.mark.parametrize( ("factory", "expected_compact"), [