Added context manager support to event subscriptions

2 files changed, 14 insertions, 15 deletions

diff --git a/src/apscheduler/schedulers/async_.py b/src/apscheduler/schedulers/async_.py
index dbece0e..213b4de 100644
--- a/src/apscheduler/schedulers/async_.py
+++ b/src/apscheduler/schedulers/async_.py
@@ -63,13 +63,14 @@ class AsyncScheduler:
 
         # Initialize the data store and start relaying events to the scheduler's event broker
         await self._exit_stack.enter_async_context(self.data_store)
-        relay_subscription = self.data_store.events.subscribe(self._events.publish)
-        self._exit_stack.callback(relay_subscription.unsubscribe)
+        self._exit_stack.enter_context(self.data_store.events.subscribe(self._events.publish))
 
         # Wake up the scheduler if the data store emits a significant schedule event
-        wakeup_subscription = self.data_store.events.subscribe(
-            lambda event: self._wakeup_event.set(), {ScheduleAdded, ScheduleUpdated})
-        self._exit_stack.callback(wakeup_subscription.unsubscribe)
+        self._exit_stack.enter_context(
+            self.data_store.events.subscribe(
+                lambda event: self._wakeup_event.set(), {ScheduleAdded, ScheduleUpdated}
+            )
+        )
 
         # Start the built-in worker, if configured to do so
         if self.start_worker:
diff --git a/src/apscheduler/schedulers/sync.py b/src/apscheduler/schedulers/sync.py
index d8e1397..dd3f37e 100644
--- a/src/apscheduler/schedulers/sync.py
+++ b/src/apscheduler/schedulers/sync.py
@@ -57,13 +57,14 @@ class Scheduler:
 
         # Initialize the data store and start relaying events to the scheduler's event broker
         self._exit_stack.enter_context(self.data_store)
-        relay_subscription = self.data_store.events.subscribe(self._events.publish)
-        self._exit_stack.callback(relay_subscription.unsubscribe)
+        self._exit_stack.enter_context(self.data_store.events.subscribe(self._events.publish))
 
         # Wake up the scheduler if the data store emits a significant schedule event
-        wakeup_subscription = self.data_store.events.subscribe(
-            lambda event: self._wakeup_event.set(), {ScheduleAdded, ScheduleUpdated})
-        self._exit_stack.callback(wakeup_subscription.unsubscribe)
+        self._exit_stack.enter_context(
+            self.data_store.events.subscribe(
+                lambda event: self._wakeup_event.set(), {ScheduleAdded, ScheduleUpdated}
+            )
+        )
 
         # Start the built-in worker, if configured to do so
         if self.start_worker:
@@ -72,12 +73,9 @@ class Scheduler:
 
         # Start the scheduler and return when it has signalled readiness or raised an exception
         start_future: Future[Event] = Future()
-        start_subscription = self._events.subscribe(start_future.set_result)
-        run_future = self._executor.submit(self.run)
-        try:
+        with self._events.subscribe(start_future.set_result):
+            run_future = self._executor.submit(self.run)
             wait([start_future, run_future], return_when=FIRST_COMPLETED)
-        finally:
-            start_subscription.unsubscribe()
 
         if run_future.done():
             run_future.result()