Add a package-level activate context manager

I believe that this has more staying power than the existing API. It could be that even this can eventually be eliminated, but I'm not sure about that just yet. Either way, this avoids exposing QueueIO as a public interface. Which is good, because I'm coming to think that the functions it provides probably all would better live in other places.

Author: ryanhiebert

README.md

@@ -21,7 +21,7 @@ Create your routines:
 # basic.py
 from time import sleep as time_sleep
 
-from queueio import QueueIO
+from queueio import activate
 from queueio import routine
 from queueio.gather import gather
 from queueio.sleep import sleep
@@ -43,11 +43,8 @@ async def yielding(iterations: int):
 
 
 if __name__ == "__main__":
-    q = QueueIO()
-    try:
-        q.submit(yielding(7))
-    finally:
-        q.shutdown()
+    with activate():
+        yielding(7).start()
 
 ```
 

queueio/__init__.py

@@ -1,2 +1,10 @@
-from .queueio import QueueIO as QueueIO
+from contextlib import contextmanager
+
+from .queueio import QueueIO as RealQueueIO
 from .registry import routine as routine
+
+
+@contextmanager
+def activate():
+    with RealQueueIO().activate():
+        yield

queueio/queueio.md

@@ -5,7 +5,9 @@
 from collections.abc import Iterable
 from concurrent.futures import Future
 from contextlib import contextmanager
+from contextvars import ContextVar
 from pathlib import Path
+from typing import Self
 
 from .broker import Broker
 from .consumer import Consumer
@@ -24,6 +26,15 @@
 
 
 class QueueIO:
+    __active = ContextVar[Self | None]("active", default=None)
+
+    @classmethod
+    def active(cls) -> Self:
+        """Find the currently active instance."""
+        queueio = cls.__active.get()
+        assert queueio is not None, "No active QueueIO instance"
+        return queueio
+
     def __init__(
         self,
         *,
@@ -35,6 +46,16 @@ def __init__(
         self.__invocations = dict[Invocation, Message]()
         self.__register_routines()
 
+    @contextmanager
+    def activate(self):
+        token = self.__active.set(self)
+        try:
+            with self.invocation_handler():
+                yield
+        finally:
+            self.__active.reset(token)
+            self.shutdown()
+
     def __pyproject(self) -> Path | None:
         for path in [cwd := Path.cwd(), *cwd.parents]:
             candidate = path / "pyproject.toml"

queueio/queueio.py

@@ -2,7 +2,7 @@
 
 from time import sleep as time_sleep
 
-from queueio import QueueIO
+from queueio import activate
 from queueio import routine
 from queueio.gather import gather
 from queueio.sleep import sleep
@@ -24,8 +24,5 @@ async def yielding(iterations: int):
 
 
 if __name__ == "__main__":
-    q = QueueIO()
-    try:
-        q.submit(yielding(7))
-    finally:
-        q.shutdown()
+    with activate():
+        yielding(7).start()

queueio/samples/basic.py

@@ -3,8 +3,8 @@
 
 import pytest
 
-from queueio import QueueIO
 from queueio.invocation import Invocation
+from queueio.queueio import QueueIO
 
 from .basic import yielding
 
@@ -13,7 +13,7 @@
 def test_integration():
     queueio = QueueIO()
 
-    try:
+    with queueio.activate():
         queueio.purge(queue="queueio")
         events = queueio.subscribe({Invocation.Completed})
         invocation = yielding(7)
@@ -36,6 +36,3 @@ def test_integration():
                 except subprocess.TimeoutExpired:
                     proc.kill()
                     proc.wait()
-
-    finally:
-        queueio.shutdown()

queueio/samples/basic_test.py

@@ -1,7 +1,7 @@
 from contextlib import suppress
 from time import sleep as time_sleep
 
-from queueio import QueueIO
+from queueio import activate
 from queueio import routine
 from queueio.gather import gather
 from queueio.sleep import sleep
@@ -47,8 +47,5 @@ async def irregular():
 
 
 if __name__ == "__main__":
-    q = QueueIO()
-    try:
-        q.submit(irregular())
-    finally:
-        q.shutdown()
+    with activate():
+        irregular().start()

queueio/samples/expanded.py

@@ -3,8 +3,8 @@
 
 import pytest
 
-from queueio import QueueIO
 from queueio.invocation import Invocation
+from queueio.queueio import QueueIO
 
 from .expanded import irregular
 
@@ -14,7 +14,7 @@ def test_integration():
     # Prefers a clean environment and queue
     queueio = QueueIO()
 
-    try:
+    with queueio.activate():
         queueio.purge(queue="queueio")
         events = queueio.subscribe({Invocation.Completed})
         invocation = irregular()
@@ -39,7 +39,3 @@ def test_integration():
                     # Force kill if it doesn't terminate gracefully
                     proc.kill()
                     proc.wait()
-
-    finally:
-        # Always clean up the queueio instance
-        queueio.shutdown()