ltq 0.1.1__tar.gz → 0.1.2__tar.gz

{ltq-0.1.1 → ltq-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: ltq
-Version: 0.1.1
+Version: 0.1.2
 Summary: Add your description here
 Author: Tom Clesius
 Author-email: Tom Clesius <tomclesius@gmail.com>

{ltq-0.1.1 → ltq-0.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltq"
-version = "0.1.1"
+version = "0.1.2"
 description = "Add your description here"
 readme = "README.md"
 authors = [{ name = "Tom Clesius", email = "tomclesius@gmail.com" }]

{ltq-0.1.1 → ltq-0.1.2}/src/ltq/task.py

@@ -1,18 +1,20 @@
 from __future__ import annotations
 
-from typing import Any, Awaitable, Callable, Generic, ParamSpec
+from functools import update_wrapper
+from typing import Awaitable, Callable, Generic, ParamSpec, TypeVar
 
 from .message import Message
 from .q import Queue
 
 P = ParamSpec("P")
+R = TypeVar("R")
 
 
-class Task(Generic[P]):
+class Task(Generic[P, R]):
     def __init__(
         self,
         name: str,
-        fn: Callable[P, Awaitable[Any]],
+        fn: Callable[P, Awaitable[R]],
         queue: Queue,
         ttl: int | None = None,
     ) -> None:
@@ -36,3 +38,6 @@ class Task(Generic[P]):
     async def send_bulk(self, messages: list[Message]) -> list[str]:
         await self.queue.put(messages, ttl=self.ttl)
         return [message.id for message in messages]
+
+    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        return await self.fn(*args, **kwargs)

{ltq-0.1.1 → ltq-0.1.2}/src/ltq/worker.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import asyncio
 from functools import partial
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, ParamSpec
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, ParamSpec, TypeVar
 
 import redis.asyncio as redis
 
@@ -21,6 +21,7 @@ logger = get_logger()
 
 
 P = ParamSpec("P")
+R = TypeVar("R")
 
 
 class Worker:
@@ -41,8 +42,8 @@ class Worker:
         self,
         queue_name: str | None = None,
         ttl: int | None = None,
-    ) -> Callable[[Callable[P, Awaitable[Any]]], Task[P]]:
-        def decorator(fn: Callable[P, Awaitable[Any]]) -> Task[P]:
+    ) -> Callable[[Callable[P, Awaitable[R]]], Task[P, R]]:
+        def decorator(fn: Callable[P, Awaitable[R]]) -> Task[P, R]:
             filename = Path(fn.__code__.co_filename).stem
             task_name = f"{filename}:{fn.__qualname__}"
             queue = Queue(self.client, queue_name or task_name)
@@ -50,7 +51,7 @@ class Worker:
                 name=task_name,
                 fn=fn,
                 queue=queue,
-                ttl=ttl, 
+                ttl=ttl,
             )
             self.tasks.append(task)
             return task
@@ -89,5 +90,8 @@ class Worker:
             await task.queue.ack(messages)
 
     async def run(self) -> None:
-        workers = (self.worker(task) for task in self.tasks)
-        await asyncio.gather(*workers)
+        try:
+            workers = (self.worker(task) for task in self.tasks)
+            await asyncio.gather(*workers)
+        finally:
+            await self.client.aclose()