feat(psi_device_base): add default timeout argument and related functionality

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:

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):

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:

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:

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

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"),
     [