Add support for watchers into AsyncStatus.wrap

[abbiemery]

Pass is as an extra argument into the wrapped function.

[coretl]

An idea for the signature:

T = TypeVar('T')
P = ParamSpec('P')
Watchers = Sequence[Callable]

class AsyncStatus(Status):
...
    @classmethod
    def wrap(cls, f: Callable[Concatenate[T, Watchers, P], Coroutine]) -> Callable[Concatenate[T, P], "AsyncStatus"]:
        @functools.wraps(f)
        def wrap_f(self: T, *args, **kwargs) -> AsyncStatus:
            watchers = []
            return AsyncStatus(f(self, watchers, *args, **kwargs), watchers)

        return wrap_f

Then we can use it to wrap prepare and set as well as kickoff and complete

[coretl]

Thinking more about this, I wonder if we should support signatures with and without watchers, i.e.

T = TypeVar('T')
P = ParamSpec('P')
Watchers = Sequence[Callable]

class AsyncStatus(Status):
...
    @classmethod
    def wrap(cls, f: Callable[P, Awaitable]) -> Callable[P, "AsyncStatus"]:
        @functools.wraps(f)
        def wrap_f(*args, **kwargs) -> AsyncStatus:
            return AsyncStatus(f(*args, **kwargs))

        return wrap_f

    @classmethod
    def wrap_passing_watchers(cls, f: Callable[Concatenate[T, Watchers, P], Awaitable]) -> Callable[Concatenate[T, P], "AsyncStatus"]:
        @functools.wraps(f)
        def wrap_f(self: T, *args, **kwargs) -> AsyncStatus:
            watchers = []
            return AsyncStatus(f(self, watchers, *args, **kwargs), watchers)

        return wrap_f

[coretl]

This would mean we could do:

class MyDevice(Stageable, Movable):
    @AsyncStatus.wrap
    def stage(self):
        ...
    @AsyncStatus.wrap_passing_watchers
    def set(self, watchers: Watchers, value: float):
        ...

Rather than passing a superfluous watchers to stage(). There is inspect magic we could do to put this in the same wrap decorator but I'm not sure it's worth it...

[coretl]

@DominicOram @dperl-dls bike shedding welcome...

[dperl-dls]

I agree with you, I'd prefer two separate wrappers for that. I must admit I don't understand how watchers is passed along in that case? It looks like it can only be an empty list? is it meant to be in the args for wrap_f?

[coretl]

The flow looks like this:

• bluesky calls st = device.set(value) which kicks off the motor moving an monitors its readback value
• some code that is interested calls st.watch(callback) which adds it to the watchers list
• on update of a motor value, the ophyd async calls for watcher in watchers: watcher(...)

Thinking about this now, this means that there is a race condition between the first update coming in and the interested code calling st.watch. We could avoid this if we instead passed an update_watchers function into the wrapped coroutine, this could then do the calling back of the added watcher with the last watched value.

[coretl]

Here's another idea, we make the wrapped function either return a coroutine or an async iterator, where the yield value is a dataclass of watcher values. Something like:

@AsyncStatus.wrap
async def stage(self) -> None:
    await self.stop_.trigger()

@AsyncStatus.wrap
async def set(self, value) -> AsyncIterator[WatcherUpdate]:
    await self.setpoint.set(value)
    for readback in observe_value(self.readback):
         yield WatcherUpdate(target=value, current=value, ...)

[dperl-dls]

oh, so that the wrapped coroutine has the same watchers as the Status, right, thanks.

[dperl-dls]

do we always want set to be doing that? there's probably not much harm in it. I like that the Status then gets updates about where it's at, that could be very useful

[coretl]

do we always want set to be doing that? there's probably not much harm in it. I like that the Status then gets updates about where it's at, that could be very useful

Yes, we do this for motors at the moment, and will do that with flyers too:

ophyd-async/src/ophyd_async/epics/demo/__init__.py

Lines 82 to 95 in fd6b8ab

	await self.setpoint.set(new_position, wait=False)
	async for current_position in observe_value(self.readback):
	for watcher in watchers:
	watcher(
	name=self.name,
	current=current_position,
	initial=old_position,
	target=new_position,
	unit=units,
	precision=precision,
	time_elapsed=time.monotonic() - start,
	)
	if np.isclose(current_position, new_position):
	break

But switching to yielding a dataclass would be a bit more type safe than calling a random callable in many places.

oh, so that the wrapped coroutine has the same watchers as the Status, right, thanks.

Thinking about this more, maybe we want to separate the watch() method into a separate class a bit like ophyd does. Then we would have:

class AsyncStatus(Status):
    def __init__(self, awaitable: Awaitable):
        # Spawn a task as before, but don't make a watcher list

    @classmethod
    def wrap(cls: Type[T], f: Callable[P, Awaitable]) -> Callable[P, T]:
        @functools.wraps(f)
        def wrap_f(*args, **kwargs) -> AsyncStatus:
            return cls(f(*args, **kwargs))

        return wrap_f

class WatchableAsyncStatus(Status):
    def __init__(self, iterator: AsyncIterator[WatcherUpdate]):
        self._watchers: List[Callable] = []
        self._start = time.monotonic()
        self._last_update: Optional[WatcherUpdate] = None
        super().__init__(self._notify_watchers_from(iterator))

    async def _notify_watchers_from(self, iterator: AsyncIterator[WatcherUpdate]):
        async for self._last_update in iterator:
            for watcher in self._watchers:
               self._update_watcher(watcher, self_last_update)

    def _update_watcher(self, watcher: Callable, update: WatcherUpdate):
        watcher(
            name=update.name,
            current=update.current,
            initial=update.initial,
            target=update.target,
            unit=update.units,
            precision=update.precision,
            time_elapsed=time.monotonic() - self._start,
        )

    def watch(self, watcher: Callable):
        self._watchers.append(watcher)
        if self._last_update:
            self._update_watcher(watcher, self._last_update)
        
    @classmethod
    def wrap(cls: Type[T], f: Callable[P, AsyncIterator[WatcherUpdate]]) -> Callable[P, T]:
        @functools.wraps(f)
        def wrap_f(*args, **kwargs) -> AsyncStatus:
            return cls(f(*args, **kwargs))

        return wrap_f

Then we get to write:

@AsyncStatus.wrap
async def stage(self) -> None:
    await self.stop_.trigger()

@WatchableAsyncStatus.wrap
async def set(self, value) -> AsyncIterator[WatcherUpdate]:
    await self.setpoint.set(value)
    for readback in observe_value(self.readback):
         yield WatcherUpdate(target=value, current=value, ...)

This also solves the first update race condition