Add custom timeout exception to be raised upon timeout in status.wait(timeout=2...)

ophyd/status.py

@@ -439,7 +439,7 @@ def exception(self, timeout=None):
         return self._exception
 
     @tracer.start_as_current_span(f"{_TRACE_PREFIX} wait")
-    def wait(self, timeout=None):
+    def wait(self, timeout=None, timeout_exception=None):
         """
         Block until the action completes.
 
@@ -451,6 +451,11 @@ def wait(self, timeout=None):
         timeout: Union[Number, None], optional
             If None (default) wait indefinitely until the status finishes.
 
+        timeout_exception: Exception, optional
+            If None (default), WaitTimeoutError will be raised if the status
+            does not complete within ``timeout``. If this is set to an Exception,
+            that exception will be raised instead.
+
         Raises
         ------
         WaitTimeoutError
@@ -465,9 +470,14 @@ def wait(self, timeout=None):
             indicates that the action itself raised ``TimeoutError``, distinct
             from ``WaitTimeoutError`` above.
         """
+        exception_to_raise = (
+            timeout_exception
+            if isinstance(timeout_exception, Exception)
+            else WaitTimeoutError(f"Status {self!r} has not completed yet.")
+        )
         _set_trace_attributes(trace.get_current_span(), self._trace_attributes)
         if not self._event.wait(timeout=timeout):
-            raise WaitTimeoutError(f"Status {self!r} has not completed yet.")
+            raise exception_to_raise
         if self._exception is not None:
             raise self._exception
 

ophyd/tests/test_status.py

@@ -437,6 +437,15 @@ def test_wait_timeout():
     with pytest.raises(WaitTimeoutError):
         st.exception(0.01)
 
+    custom_exception = ValueError("My custom exception")
+    with pytest.raises(ValueError) as exc:
+        st.wait(0.01, timeout_exception=custom_exception)
+        assert exc.value is custom_exception
+
+    # Wrong input should raise WaitTimeoutError
+    with pytest.raises(WaitTimeoutError):
+        st.wait(0.01, timeout_exception="Not an exception")
+
 
 def test_status_timeout():
     """