Enhanced typing of subscribe()

[peterschutt]

This PR adds type var T_Event and uses it to type the subscribe() interfaces.

Fixes an issue where subscribing a handler would raise a type error if the handler is typed to receive a subclass of Event.

Allows type checkers to verify that the handler passed to subscribe() can support the event types it is assigned to handle.

Closes #846

[coveralls]

coverage: 92.772% (+0.2%) from 92.535%
when pulling 9207f3e on peterschutt:846-subscribe-static-typing
into b5dfbed on agronholm:master.

[agronholm]

I have some doubts whether this will work with multiple event types (ie. unions). Have you checked that?

[peterschutt]

TLDR;

+1:

• change improves simple case of single event type subscription (no type narrowing required in callback).

+0:

• multiple event type subscription is not any worse than status quo in behavior (still need to perform type narrowing in callback)

-1:

• inconsistent inference of event type unions
• how long it took to investigate and explain this

So I suppose that leaves me -1 (for now, I will spend some time to see if I can improve this).

---

I have some doubts whether this will work with multiple event types (ie. unions). Have you checked that?

I did, this example has no type error:

from examples import scheduler

from apscheduler import JobAcquired, JobReleased


def handler(event: JobAcquired | JobReleased) -> None:
    ...


scheduler.subscribe(handler, {JobAcquired, JobReleased})

JobAcquired and JobReleased were what I was working with at the time, however this was unfortunate because:

reveal_type({JobAcquired, JobReleased})  # Revealed type is "builtins.set[builtins.type]"

mypy narrows these down to type, not even Event. At the time I assumed there was no type error because it was type safe.

Instead, using TaskAdded and TaskUpdated b/c they are the first two in the _events module:

reveal_type({TaskAdded, TaskUpdated})
# Revealed type is "builtins.set[def (*, timestamp: Any =, task_id: builtins.str) -> apscheduler._events.DataStoreEvent]

For these, mypy reduces them to a set of callables that return a DataStoreEvent instance, so that is what the handler needs to receive:

from examples import scheduler

from apscheduler import DataStoreEvent, TaskAdded, TaskUpdated


def handler(event: DataStoreEvent) -> None:
    ...


scheduler.subscribe(handler, {TaskAdded, TaskUpdated})

B/c DataStoreEvent doesn't carry any more information than Event - this isn't any better than status quo, but its not worse either.

It would be less surprising if it would allow:

def handler(event: TaskAdded | TaskUpdated) -> None:
    ....

However, in any case, type narrowing is going to be required in a multi-event type callback.

[agronholm]

Thanks for taking the time to investigate. I'll mull it over once more.

[peterschutt]

This is the nicest I've been able to get the behavior (diff is against this branch, and only adjusts the async scheduler interface, other interfaces would have to be updated also):

diff --git a/src/apscheduler/_schedulers/async_.py b/src/apscheduler/_schedulers/async_.py
index 6c5ae94..b975c90 100644
--- a/src/apscheduler/_schedulers/async_.py
+++ b/src/apscheduler/_schedulers/async_.py
@@ -11,7 +11,7 @@ from functools import partial
 from inspect import isbuiltin, isclass, ismethod, ismodule
 from logging import Logger, getLogger
 from types import TracebackType
-from typing import Any, Callable, Iterable, Mapping, cast
+from typing import Any, Callable, Iterable, Mapping, cast, overload
 from uuid import UUID, uuid4
 
 import anyio
@@ -215,10 +215,30 @@ class AsyncScheduler:
         await self.data_store.cleanup()
         self.logger.info("Cleaned up expired job results and finished schedules")
 
+    @overload
     def subscribe(
         self,
         callback: Callable[[T_Event], Any],
-        event_types: type[T_Event] | Iterable[type[T_Event]] | None = None,
+        event_types: type[T_Event],
+        **kwargs: Any,
+    ) -> Subscription:
+        ...
+
+    @overload
+    def subscribe(
+        self,
+        callback: Callable[[Event], Any],
+        event_types: Iterable[type[Event]],
+        *,
+        one_shot: bool = ...,
+        is_async: bool = ...,
+    ) -> Subscription:
+        ...
+
+    def subscribe(
+        self,
+        callback: Callable[[T_Event], Any] | Callable[[Event], Any],
+        event_types: type[T_Event] | Iterable[type[Event]] | None = None,
         *,
         one_shot: bool = False,
         is_async: bool = True,

With this pattern it allows the callback to share the type of a single event type subscription, for example, this type checks fine:

def handler(event: JobAcquired) -> None:
    ...

scheduler.subscribe(handler, JobAcquired)

And for the iterable of types it behaves the same as it does now, i.e., this is an error:

def handler(event: TaskAdded | TaskUpdated) -> None:
    ...

scheduler.subscribe(handler, {TaskAdded, TaskUpdated})
# error: Argument 1 to "subscribe" of "AsyncScheduler" has incompatible type "Callable[[TaskAdded | TaskUpdated], None]"; expected "Callable[[Event], Any]"  [arg-type]

This is no error:

def handler(event: Event) -> None:
    ...

scheduler.subscribe(handler, {TaskAdded, TaskUpdated})

Its not a bad middle ground as it removes the need to have a redundant type narrowing or type ignore where the handler should only receive a single type, and leaves the rest of the behavior as it is.

We could go further and define overloads for DataStoreEvent and SchedulerEvent so that this would type check OK:

def handler(event: DataStoreEvent) -> None:
    ...

scheduler.subscribe(handler, {TaskAdded, TaskUpdated})

But given that neither of those types define any attributes that are common across all of their sub types, it wouldn't be of any benefit.

[agronholm]

It sounds like it's at least an improvement over the existing typing. Any idea if the test failures are caused by these changes?

[agronholm]

Hm, looks more like an incompatibility between pytest 8 and pytest-lazyfixture.

[agronholm]

Yup, looks like: TvoroG/pytest-lazy-fixture#65

[agronholm]

Alright, finally got around to merging this. Thanks!