ophyd-async 0.14.1__tar.gz → 0.15__tar.gz

{ophyd_async-0.14.1 → ophyd_async-0.15}/.copier-answers.yml

@@ -1,5 +1,5 @@
 # Changes here will be overwritten by Copier
-_commit: 5.0.1
+_commit: 5.0.2
 _src_path: https://github.com/DiamondLightSource/python-copier-template
 author_email: tom.cobb@diamond.ac.uk
 author_name: Tom Cobb

{ophyd_async-0.14.1 → ophyd_async-0.15}/.github/CONTRIBUTING.md

@@ -24,4 +24,4 @@ It is recommended that developers use a [vscode devcontainer](https://code.visua
 
 This project was created using the [Diamond Light Source Copier Template](https://github.com/DiamondLightSource/python-copier-template) for Python projects.
 
-For more information on common tasks like setting up a developer environment, running the tests, and setting a pre-commit hook, see the template's [How-to guides](https://diamondlightsource.github.io/python-copier-template/5.0.1/how-to.html).
+For more information on common tasks like setting up a developer environment, running the tests, and setting a pre-commit hook, see the template's [How-to guides](https://diamondlightsource.github.io/python-copier-template/5.0.2/how-to.html).

{ophyd_async-0.14.1/src/ophyd_async.egg-info → ophyd_async-0.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ophyd-async
-Version: 0.14.1
+Version: 0.15
 Summary: Asynchronous Bluesky hardware abstraction code, compatible with control systems like EPICS and Tango
 Author-email: Tom Cobb <tom.cobb@diamond.ac.uk>
 License: BSD 3-Clause License

{ophyd_async-0.14.1 → ophyd_async-0.15}/docs/conf.py

@@ -147,6 +147,7 @@ obj_ignore = [
     "ophyd_async.core._device_filler.SignalBackendT",
     "ophyd_async.core._device_filler.DeviceConnectorT",
     "ophyd_async.core._derived_signal_backend.TransformT",
+    "ophyd_async.core._mock_signal_backend.MockPutCallback",
     "ophyd_async.core._protocol.C",
     "ophyd_async.core._signal_backend.SignalDatatypeV",
     "ophyd_async.core._status.AsyncStatusBase",
@@ -164,8 +165,10 @@ obj_ignore = [
     "ophyd_async.testing._utils.T",
     "ophyd_async.sim._mirror.TwoJackRaw",
     "ophyd_async.sim._mirror.TwoJackDerived",
+    "non_zero",
     "0.1",
     "1.0",
+    "bluesky.protocols.T_co",
 ]
 nitpick_ignore = []
 for var in obj_ignore:

{ophyd_async-0.14.1 → ophyd_async-0.15}/docs/how-to/interact-with-signals.md

@@ -43,16 +43,9 @@ To observe every value change and run a function on that value you can use [](#o
 async for value in observe_value(signal):
     do_something_with(value)
 ```
-This will run until you `break` out of the loop. If you would like to break out of the loop when some other operation is complete you can pass an [](#AsyncStatus) that you create yourself or from the result of setting a signal, which will break out of the loop:
+This will run until you `break` out of the loop.
 
-```python
-status = signal.set(value)
-async for value in observe_value(signal2, done_status=status):
-    do_something_with(value)
-    # when signal.set() completes the loop will break out here
-```
-
-You can pass `timeout` to specify how the maximum time to wait for a single update, and `done_timeout` to specify the maximum time to wait for `done_status`.
+You can pass `timeout` to specify the maximum time to wait for a single update.
 
 If you want to wait for multiple signals you can use [](#observe_signals_value):
 ```python
@@ -63,6 +56,38 @@ async for signal, value in observe_value(signal1, signal2):
         do_something_else_with(value)
 ```
 
+## Use AsyncStatus as a context manager to bound loop execution
+
+If you want a loop to run until some operation completes, you can use [](#AsyncStatus) as a context manager. When the status completes, it will cancel the calling task, causing the loop to exit. This is useful when you want to process signal updates until another operation finishes:
+
+```python
+# Process updates while a motor is moving
+async with motor.set(target_position):
+    async for value in observe_value(detector):
+        process_reading(value)
+        # Loop automatically exits when motor reaches position
+```
+
+If the loop completes before the status, the status task is automatically cancelled:
+
+```python
+async with signal1.set(new_value):
+    for i in range(3):
+        value = await signal.get_value()
+        process(value)
+        # Loop completes after 3 iterations, cancelling the wait for signal1 to finishe being set
+```
+
+If an exception is raised in the loop body, it propagates out normally:
+
+```python
+async with signal1.set(new_value):
+    async for value in observe_value(signal2):
+        if value > threshold:
+            raise ValueError("Threshold exceeded")
+        # Exception propagates, status is cancelled and no longer waits for signal1 to finish being set
+```
+
 ## Wait for the value to match some expected value
 
 If you don't need to run code for every signal update, but just want to wait until the signal matches some expected value, you can use [](#wait_for_value):

{ophyd_async-0.14.1 → ophyd_async-0.15}/pyproject.toml

@@ -147,7 +147,7 @@ commands = [
     [
         "pyright",
         "--pythonpath",
-        ".venv/bin/python",
+        "{env:VIRTUAL_ENV}/bin/python",
         "src",
         { replace = "posargs", default = [
         ], extend = true },

{ophyd_async-0.14.1 → ophyd_async-0.15}/src/ophyd_async/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.14.1'
-__version_tuple__ = version_tuple = (0, 14, 1)
+__version__ = version = '0.15'
+__version_tuple__ = version_tuple = (0, 15)
 
-__commit_id__ = commit_id = 'g9b567b214'
+__commit_id__ = commit_id = 'g6ebd5ae5e'

{ophyd_async-0.14.1 → ophyd_async-0.15}/src/ophyd_async/core/__init__.py

@@ -1,5 +1,6 @@
 """The building blocks for making devices."""
 
+from ._datatypes import Array1D, DTypeScalar_co, Table, TableSubclass
 from ._derived_signal import (
     DerivedSignalFactory,
     derived_signal_r,
@@ -23,7 +24,7 @@ from ._device import (
     default_mock_class,
     init_devices,
 )
-from ._device_filler import DeviceFiller
+from ._device_filler import DeviceAnnotation, DeviceFiller
 from ._enums import (
     EnabledDisabled,
     EnableDisable,
@@ -44,7 +45,15 @@ from ._mock_signal_utils import (
     set_mock_value,
     set_mock_values,
 )
-from ._protocol import AsyncConfigurable, AsyncReadable, AsyncStageable, Watcher
+from ._protocol import (
+    AsyncConfigurable,
+    AsyncLocatable,
+    AsyncMovable,
+    AsyncPausable,
+    AsyncReadable,
+    AsyncStageable,
+    Watcher,
+)
 from ._providers import (
     AutoIncrementFilenameProvider,
     AutoIncrementingPathProvider,
@@ -86,8 +95,6 @@ from ._signal import (
     walk_signal_sources,
 )
 from ._signal_backend import (
-    Array1D,
-    DTypeScalar_co,
     Primitive,
     SignalBackend,
     SignalDatatype,
@@ -97,7 +104,6 @@ from ._signal_backend import (
 )
 from ._soft_signal_backend import SoftSignalBackend
 from ._status import AsyncStatus, WatchableAsyncStatus, completed_status
-from ._table import Table, TableSubclass
 from ._utils import (
     CALCULATE_TIMEOUT,
     DEFAULT_TIMEOUT,
@@ -117,6 +123,7 @@ from ._utils import (
     get_enum_cls,
     get_unique,
     in_micros,
+    non_zero,
     wait_for_connection,
 )
 from ._yaml_settings import YamlSettingsProvider
@@ -146,11 +153,15 @@ __all__ = [
     "Device",
     "DeviceConnector",
     "DeviceFiller",
+    "DeviceAnnotation",
     "DeviceVector",
     "init_devices",
     # Protocols
     "AsyncReadable",
     "AsyncConfigurable",
+    "AsyncLocatable",
+    "AsyncMovable",
+    "AsyncPausable",
     "AsyncStageable",
     "Watcher",
     # Status
@@ -255,6 +266,7 @@ __all__ = [
     "make_datakey",
     "wait_for_connection",
     "Ignore",
+    "non_zero",
     # Derived signal
     "derived_signal_r",
     "derived_signal_rw",

ophyd_async-0.14.1/src/ophyd_async/core/_table.py → ophyd_async-0.15/src/ophyd_async/core/_datatypes.py

@@ -1,13 +1,25 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Sequence
-from typing import Annotated, Any, TypeVar, get_origin, get_type_hints
+from typing import Annotated, Any, TypeVar
 
 import numpy as np
 from pydantic import ConfigDict, Field, model_validator
 from pydantic_numpy.helper.annotation import NpArrayPydanticAnnotation
 
-from ._utils import ConfinedModel, get_dtype
+from ._utils import ConfinedModel, cached_get_origin, cached_get_type_hints, get_dtype
+
+DTypeScalar_co = TypeVar("DTypeScalar_co", covariant=True, bound=np.generic)
+"""A numpy dtype like [](#numpy.float64)."""
+
+
+# To be a 1D array shape should really be tuple[int], but np.array()
+# currently produces tuple[int, ...] even when it has 1D input args
+# https://github.com/numpy/numpy/issues/28077#issuecomment-2566485178
+Array1D = np.ndarray[tuple[int, ...], np.dtype[DTypeScalar_co]]
+"""A type alias for a 1D numpy array with a specific scalar data type.
+
+E.g. `Array1D[np.float64]` is a 1D numpy array of 64-bit floats."""
 
 TableSubclass = TypeVar("TableSubclass", bound="Table")
 
@@ -75,11 +87,8 @@ class Table(ConfinedModel):
         # ...but forbid extra in subclasses so it gets validated
         cls.model_config = ConfigDict(validate_assignment=True, extra="forbid")
         # Change fields to have the correct annotations
-        # TODO: refactor so we don't need this to break circular imports
-        from ._signal_backend import Array1D
-
-        for k, anno in get_type_hints(cls, localns={"Array1D": Array1D}).items():
-            if get_origin(anno) is np.ndarray:
+        for k, anno in cached_get_type_hints(cls).items():
+            if cached_get_origin(anno) is np.ndarray:
                 dtype = get_dtype(anno)
                 new_anno = Annotated[
                     anno,
@@ -88,7 +97,7 @@ class Table(ConfinedModel):
                     ),
                     Field(default_factory=_make_default_factory(dtype)),
                 ]
-            elif get_origin(anno) is Sequence:
+            elif cached_get_origin(anno) is Sequence:
                 new_anno = Annotated[anno, Field(default_factory=list)]
             else:
                 raise TypeError(f"Cannot use annotation {anno} in a Table")
@@ -156,7 +165,7 @@ class Table(ConfinedModel):
             raise AssertionError(f"Cannot construct Table from {data}")
         for field_name, field_value in cls.model_fields.items():
             if (
-                get_origin(field_value.annotation) is np.ndarray
+                cached_get_origin(field_value.annotation) is np.ndarray
                 and field_value.annotation
                 and field_name in data_dict
             ):

{ophyd_async-0.14.1 → ophyd_async-0.15}/src/ophyd_async/core/_derived_signal.py

@@ -1,5 +1,14 @@
-from collections.abc import Awaitable, Callable
-from typing import Any, Generic, get_args, get_origin, get_type_hints, is_typeddict
+import functools
+from collections.abc import Awaitable, Callable, Mapping
+from inspect import Parameter, signature
+from typing import (
+    Any,
+    Generic,
+    TypeVar,
+    get_args,
+    get_origin,
+    is_typeddict,
+)
 
 from bluesky.protocols import Locatable
 
@@ -12,6 +21,7 @@ from ._derived_signal_backend import (
 from ._device import Device
 from ._signal import Signal, SignalR, SignalRW, SignalT, SignalW
 from ._signal_backend import Primitive, SignalDatatypeT
+from ._utils import cached_get_type_hints
 
 
 class DerivedSignalFactory(Generic[TransformT]):
@@ -53,12 +63,13 @@ class DerivedSignalFactory(Generic[TransformT]):
             # Populate expected parameters and types
             expected = {
                 **{k: f.annotation for k, f in transform_cls.model_fields.items()},
-                **{
-                    k: v
-                    for k, v in get_type_hints(transform_cls.raw_to_derived).items()
-                    if k not in {"self", "return"}
-                },
+                **_get_params_types_dict(transform_cls.raw_to_derived),
             }
+            if empty_keys := [k for k, v in expected.items() if v == Parameter.empty]:
+                raise TypeError(
+                    f"{transform_cls.raw_to_derived} is missing a type "
+                    f"hint for arguments: {empty_keys}"
+                )
 
             # Populate received parameters and types
             # Use Primitive's type, Signal's datatype,
@@ -76,7 +87,19 @@ class DerivedSignalFactory(Generic[TransformT]):
                     f"Expected the following to be passed as keyword arguments "
                     f"{expected}, got {received}"
                 )
-                raise TypeError(msg)
+                if set(expected.keys()) - set(received.keys()):
+                    raise TypeError(msg)
+
+                for k in set(expected.keys()):
+                    if isinstance(expected[k], type):
+                        if not issubclass(received[k], expected[k]):
+                            raise TypeError(msg)
+                    elif isinstance(expected[k], TypeVar):
+                        bound = expected[k].__bound__
+                        if isinstance(bound, type) and not issubclass(
+                            received[k], bound
+                        ):
+                            raise TypeError(msg)
         self._set_derived_takes_dict = (
             is_typeddict(_get_first_arg_datatype(set_derived)) if set_derived else False
         )
@@ -185,7 +208,7 @@ class DerivedSignalFactory(Generic[TransformT]):
 
 
 def _get_return_datatype(func: Callable[..., SignalDatatypeT]) -> type[SignalDatatypeT]:
-    args = get_type_hints(func)
+    args = cached_get_type_hints(func)
     if "return" not in args:
         msg = f"{func} does not have a type hint for it's return value"
         raise TypeError(msg)
@@ -195,28 +218,28 @@ def _get_return_datatype(func: Callable[..., SignalDatatypeT]) -> type[SignalDat
 def _get_first_arg_datatype(
     func: Callable[[SignalDatatypeT], Any],
 ) -> type[SignalDatatypeT]:
-    args = get_type_hints(func)
-    args.pop("return", None)
+    args = _get_params_types_dict(func)
     if not args:
         msg = f"{func} does not have a type hinted argument"
         raise TypeError(msg)
     return list(args.values())[0]
 
 
+def _get_params_types_dict(inspected_function: Callable) -> Mapping[str, Any]:
+    sig = signature(inspected_function, eval_str=True)
+    exclude_keys = {"self", "args", "kwargs", "cls"}
+    return {k: v.annotation for k, v in sig.parameters.items() if k not in exclude_keys}
+
+
 def _make_factory(
-    raw_to_derived: Callable[..., SignalDatatypeT] | None = None,
+    raw_to_derived_func: Callable[..., SignalDatatypeT] | None = None,
     set_derived: Callable[[SignalDatatypeT], Awaitable[None]] | None = None,
     raw_devices_and_constants: dict[str, Device | Primitive] | None = None,
 ) -> DerivedSignalFactory:
-    if raw_to_derived:
+    if raw_to_derived_func:
 
         class DerivedTransform(Transform):
-            def raw_to_derived(self, **kwargs) -> dict[str, SignalDatatypeT]:
-                return {"value": raw_to_derived(**kwargs)}
-
-        # Update the signature for raw_to_derived to match what we are passed as this
-        # will be checked in DerivedSignalFactory
-        DerivedTransform.raw_to_derived.__annotations__ = get_type_hints(raw_to_derived)
+            raw_to_derived = _dict_wrapper(raw_to_derived_func)
 
         return DerivedSignalFactory(
             DerivedTransform,
@@ -245,7 +268,7 @@ def derived_signal_r(
         The names of these arguments must match the arguments of raw_to_derived.
     """
     factory = _make_factory(
-        raw_to_derived=raw_to_derived,
+        raw_to_derived_func=raw_to_derived,
         raw_devices_and_constants=raw_devices_and_constants,
     )
     return factory.derived_signal_r(
@@ -278,16 +301,16 @@ def derived_signal_rw(
         The names of these arguments must match the arguments of raw_to_derived.
     """
     raw_to_derived_datatype = _get_return_datatype(raw_to_derived)
-    set_derived_datatype = _get_first_arg_datatype(set_derived)
-    if raw_to_derived_datatype != set_derived_datatype:
+    set_derived_arg_datatype = _get_first_arg_datatype(set_derived)
+    if raw_to_derived_datatype != set_derived_arg_datatype:
         msg = (
             f"{raw_to_derived} has datatype {raw_to_derived_datatype} "
-            f"!= {set_derived_datatype} datatype {set_derived_datatype}"
+            f"!= {set_derived_arg_datatype} datatype {set_derived_arg_datatype}"
         )
         raise TypeError(msg)
 
     factory = _make_factory(
-        raw_to_derived=raw_to_derived,
+        raw_to_derived_func=raw_to_derived,
         set_derived=set_derived,
         raw_devices_and_constants=raw_devices_and_constants,
     )
@@ -343,3 +366,13 @@ def _partition_by_keys(data: dict, keys: set) -> tuple[dict, dict]:
         else:
             group_excluded[k] = v
     return group_excluded, group_included
+
+
+def _dict_wrapper(
+    fn: Callable[..., SignalDatatypeT],
+) -> Callable[..., dict[str, SignalDatatypeT]]:
+    @functools.wraps(fn)
+    def wrapped(self, **kwargs):
+        return {"value": fn(**kwargs)}
+
+    return wrapped

{ophyd_async-0.14.1 → ophyd_async-0.15}/src/ophyd_async/core/_derived_signal_backend.py

@@ -291,11 +291,7 @@ class DerivedSignalBackend(SignalBackend[SignalDatatypeT]):
         )
         raise RuntimeError(msg)
 
-    async def put(self, value: SignalDatatypeT | None, wait: bool) -> None:
-        if wait is False:
-            msg = "Cannot put with wait=False"
-            raise RuntimeError(msg)
-
+    async def put(self, value: SignalDatatypeT | None) -> None:
         value = error_if_none(
             value,
             "Must be given a value to put",

{ophyd_async-0.14.1 → ophyd_async-0.15}/src/ophyd_async/core/_device_filler.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import types
 from abc import abstractmethod
 from collections.abc import Callable, Iterator, Sequence
 from typing import (
@@ -9,16 +10,16 @@ from typing import (
     NoReturn,
     Protocol,
     TypeVar,
+    Union,
     cast,
     get_args,
-    get_type_hints,
     runtime_checkable,
 )
 
 from ._device import Device, DeviceConnector, DeviceVector
 from ._signal import Ignore, Signal, SignalX
 from ._signal_backend import SignalBackend, SignalDatatype
-from ._utils import get_origin_class
+from ._utils import cached_get_origin, cached_get_type_hints, get_origin_class
 
 SignalBackendT = TypeVar("SignalBackendT", bound=SignalBackend)
 DeviceConnectorT = TypeVar("DeviceConnectorT", bound=DeviceConnector)
@@ -76,6 +77,7 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
         self._extras: dict[UniqueName, Sequence[Any]] = {}
         self._signal_datatype: dict[LogicalName, type | None] = {}
         self._vector_device_type: dict[LogicalName, type[Device] | None] = {}
+        self._optional_devices: set[str] = set()
         self.ignored_signals: set[str] = set()
         # Backends and Connectors stored ready for the connection phase
         self._unfilled_backends: dict[
@@ -113,14 +115,28 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
         # https://github.com/python/cpython/issues/124840
         cls = type(self._device)
         # Get hints without Annotated for determining types
-        hints = get_type_hints(cls)
+        hints = cached_get_type_hints(cls)
         # Get hints with Annotated for wrapping signals and backends
-        extra_hints = get_type_hints(cls, include_extras=True)
+        extra_hints = cached_get_type_hints(cls, include_extras=True)
         for attr_name, annotation in hints.items():
             if annotation is Ignore:
                 self.ignored_signals.add(attr_name)
             name = UniqueName(attr_name)
             origin = get_origin_class(annotation)
+            args = get_args(annotation)
+
+            if (
+                cached_get_origin(annotation) is Union
+                and types.NoneType in args
+                and len(args) == 2
+            ):
+                # Annotation is an Union with two arguments, one of which is None
+                # Make this signal an optional parameter and set origin to T
+                # so the device is added to unfilled_connectors
+                self._optional_devices.add(name)
+                (annotation,) = [x for x in args if x is not types.NoneType]
+                origin = get_origin_class(annotation)
+
             if (
                 name == "parent"
                 or name.startswith("_")
@@ -241,10 +257,17 @@ class DeviceFiller(Generic[SignalBackendT, DeviceConnectorT]):
         :param source: The source of the data that should have done the filling, for
                        reporting as an error message
         """
-        unfilled = sorted(set(self._unfilled_connectors).union(self._unfilled_backends))
-        if unfilled:
+        unfilled = set(self._unfilled_connectors).union(self._unfilled_backends)
+        unfilled_optional = sorted(unfilled.intersection(self._optional_devices))
+
+        for name in unfilled_optional:
+            setattr(self._device, name, None)
+
+        required = sorted(unfilled.difference(unfilled_optional))
+
+        if required:
             raise RuntimeError(
-                f"{self._device.name}: cannot provision {unfilled} from {source}"
+                f"{self._device.name}: cannot provision {required} from {source}"
             )
 
     def _ensure_device_vector(self, name: LogicalName) -> DeviceVector:

{ophyd_async-0.14.1 → ophyd_async-0.15}/src/ophyd_async/core/_mock_signal_backend.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from functools import cached_property
 from typing import TYPE_CHECKING
 from unittest.mock import AsyncMock
@@ -17,6 +17,11 @@ from ._utils import Callback
 if TYPE_CHECKING:
     from ._device import LazyMock
 
+MockPutCallback = (
+    Callable[[SignalDatatypeT], SignalDatatypeT | None]
+    | Callable[[SignalDatatypeT], Awaitable[SignalDatatypeT | None]]
+)
+
 
 class MockSignalBackend(SignalBackend[SignalDatatypeT]):
     """Signal backend for testing, created by ``Device.connect(mock=True)``."""
@@ -42,13 +47,27 @@ class MockSignalBackend(SignalBackend[SignalDatatypeT]):
 
         # use existing Mock if provided
         self.mock = mock
+        self._mock_put_callback: MockPutCallback | None = None
         super().__init__(datatype=self.initial_backend.datatype)
 
+    def set_mock_put_callback(self, callback: MockPutCallback | None):
+        if "put_mock" in self.__dict__:
+            # put_mock cached property exists, so set the side effect on it
+            self.put_mock.side_effect = callback
+        else:
+            # put_mock doesn't exist, don't create it as that would be slow
+            # so just keep it internally
+            self._mock_put_callback = callback
+
     @cached_property
     def put_mock(self) -> AsyncMock:
         """Return the mock that will track calls to `put()`."""
         put_mock = AsyncMock(
-            name="put", spec=Callable, side_effect=lambda *_, **__: None
+            name="put",
+            spec=Callable,
+            side_effect=self._mock_put_callback
+            if self._mock_put_callback
+            else lambda v: None,
         )
         self.mock().attach_mock(put_mock, "put")
         return put_mock
@@ -73,13 +92,12 @@ class MockSignalBackend(SignalBackend[SignalDatatypeT]):
         put_proceeds.set()
         return put_proceeds
 
-    async def put(self, value: SignalDatatypeT | None, wait: bool):
-        new_value = await self.put_mock(value, wait=wait)
+    async def put(self, value: SignalDatatypeT | None):
+        new_value = await self.put_mock(value)
         if new_value is None:
             new_value = value
-        await self.soft_backend.put(new_value, wait=wait)
-        if wait:
-            await self.put_proceeds.wait()
+        await self.soft_backend.put(new_value)
+        await self.put_proceeds.wait()
 
     async def get_reading(self) -> Reading:
         return await self.soft_backend.get_reading()

{ophyd_async-0.14.1 → ophyd_async-0.15}/src/ophyd_async/core/_mock_signal_utils.py

@@ -1,9 +1,9 @@
-from collections.abc import Awaitable, Callable, Iterable, Iterator
+from collections.abc import Iterable, Iterator
 from contextlib import contextmanager
 from unittest.mock import AsyncMock, Mock
 
 from ._device import Device, DeviceMock
-from ._mock_signal_backend import MockSignalBackend
+from ._mock_signal_backend import MockPutCallback, MockSignalBackend
 from ._signal import Signal, SignalConnector, SignalR
 from ._signal_backend import SignalDatatypeT
 
@@ -109,16 +109,12 @@ def set_mock_values(
 
 
 @contextmanager
-def _unset_side_effect_cm(put_mock: AsyncMock):
+def _unset_side_effect_cm(backend: MockSignalBackend):
     yield
-    put_mock.side_effect = None
+    backend.set_mock_put_callback(None)
 
 
-def callback_on_mock_put(
-    signal: Signal[SignalDatatypeT],
-    callback: Callable[[SignalDatatypeT, bool], SignalDatatypeT | None]
-    | Callable[[SignalDatatypeT, bool], Awaitable[SignalDatatypeT | None]],
-):
+def callback_on_mock_put(signal: Signal[SignalDatatypeT], callback: MockPutCallback):
     """For setting a callback when a backend is put to.
 
     Can either be used in a context, with the callback being unset on exit, or
@@ -132,8 +128,8 @@ def callback_on_mock_put(
         context.
     """
     backend = _get_mock_signal_backend(signal)
-    backend.put_mock.side_effect = callback
-    return _unset_side_effect_cm(backend.put_mock)
+    backend.set_mock_put_callback(callback)
+    return _unset_side_effect_cm(backend)
 
 
 def set_mock_put_proceeds(signal: Signal, proceeds: bool):