qena-shared-lib 0.1.10__py3-none-any.whl → 0.1.12__py3-none-any.whl

qena_shared_lib/rabbitmq/_listener.py

@@ -59,9 +59,9 @@ class FlowControl:
         self._channel = channel
         self._loop = loop
         self._lock = Lock()
-        self._flow_control_future = None
+        self._flow_control_future: Future | None = None
 
-    async def request(self, prefetch_count: int):
+    async def request(self, prefetch_count: int) -> None:
         async with self._lock:
             self._flow_control_future = self._loop.create_future()
 
@@ -72,7 +72,7 @@ class FlowControl:
 
             await self._flow_control_future
 
-    def _on_prefetch_count_set(self, method: Method):
+    def _on_prefetch_count_set(self, method: Method) -> None:
         del method
 
         if self._flow_control_future is None:
@@ -87,13 +87,13 @@ class RpcReply:
         channel_pool: ChannelPool,
         reply_to: str,
         correlation_id: str | None = None,
-    ):
+    ) -> None:
         self._channel_pool = channel_pool
         self._reply_to = reply_to
         self._correlation_id = correlation_id
         self._replied = False
 
-    async def reply(self, message: Any):
+    async def reply(self, message: Any) -> None:
         base_channel = await self._channel_pool.get()
 
         reply_properties = BasicProperties(content_type="application/json")
@@ -128,7 +128,7 @@ class ListenerContext:
     rpc_reply: RpcReply | None = None
     context_dispose_callback: Callable | None = None
 
-    def dispose(self):
+    def dispose(self) -> None:
         if self.context_dispose_callback is not None:
             self.context_dispose_callback(self)
 
@@ -145,7 +145,7 @@ class BackoffRetryDelay(RetryDelayStrategy):
     min: float
     max: float
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.min > self.max:
             raise ValueError("`min` greater than `max`")
 
@@ -173,7 +173,7 @@ class RetryDelayJitter:
     min: float = 0.5
     max: float = 1.0
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.min > self.max:
             raise ValueError("`min` greater than `max`")
 
@@ -214,7 +214,7 @@ class ListenerChannelAdapter(BaseChannel):
         connection: AsyncioConnection,
         on_channel_open_callback: Callable[[Channel], None],
         on_cancel_callback: Callable,
-    ):
+    ) -> None:
         super().__init__(
             connection=connection,
             failed_reopen_threshold=None,
@@ -223,10 +223,10 @@ class ListenerChannelAdapter(BaseChannel):
         self._on_channle_open_listener_callback = on_channel_open_callback
         self._on_listener_cancel_callback = on_cancel_callback
 
-    def _hook_on_channel_opened(self):
+    def _hook_on_channel_opened(self) -> None:
         self._on_channle_open_listener_callback(self.channel)
 
-    def _hook_on_cancelled(self):
+    def _hook_on_cancelled(self) -> None:
         self._on_listener_cancel_callback()
 
 
@@ -265,7 +265,7 @@ class Listener(AsyncEventLoopMixin):
         durable: bool = True,
         purge_on_startup: bool = False,
         retry_policy: RetryPolicy | None = None,
-    ):
+    ) -> None:
         self._queue = queue
         self._listener_name_header_key = listener_name_header_key
         self._prefetch_count = prefetch_count
@@ -292,7 +292,7 @@ class Listener(AsyncEventLoopMixin):
         listener_name: str | None,
         listener_method: Callable,
         retry_policy: RetryPolicy | None = None,
-    ):
+    ) -> None:
         self._register_listener_method(
             listener_name=listener_name,
             listener_method=listener_method,
@@ -306,7 +306,7 @@ class Listener(AsyncEventLoopMixin):
         listener_method: Callable,
         parameters: MappingProxyType[str, Parameter],
         retry_policy: RetryPolicy | None = None,
-    ):
+    ) -> None:
         listener_name = listener_name or "__default__"
 
         if listener_name in self._listeners:
@@ -345,7 +345,7 @@ class Listener(AsyncEventLoopMixin):
         container: Container,
         logstash: BaseLogstashSender,
         global_retry_policy: RetryPolicy | None = None,
-    ):
+    ) -> None:
         self._connection = connection
         self._channel_pool = channel_pool
         self._listener_future = self.loop.create_future()
@@ -362,16 +362,16 @@ class Listener(AsyncEventLoopMixin):
 
         await self._listener_future
 
-    def _on_channel_opened(self, channel: Channel):
+    def _on_channel_opened(self, channel: Channel) -> None:
         self._channel = channel
         self._flow_control = FlowControl(channel=self._channel, loop=self.loop)
 
         self._declare_queue()
 
-    def _on_cancelled(self):
+    def _on_cancelled(self) -> None:
         self._declare_queue()
 
-    def _declare_queue(self):
+    def _declare_queue(self) -> None:
         try:
             self._channel.queue_declare(
                 queue=self._queue,
@@ -381,7 +381,7 @@ class Listener(AsyncEventLoopMixin):
         except Exception as e:
             self._fail_listener(e)
 
-    def _on_queue_declared(self, method: Method):
+    def _on_queue_declared(self, method: Method) -> None:
         del method
 
         if self._purge_on_startup:
@@ -395,7 +395,7 @@ class Listener(AsyncEventLoopMixin):
         else:
             self._set_prefetch_count()
 
-    def _set_prefetch_count(self):
+    def _set_prefetch_count(self) -> None:
         try:
             self._channel.basic_qos(
                 prefetch_count=self._prefetch_count,
@@ -404,7 +404,7 @@ class Listener(AsyncEventLoopMixin):
         except Exception as e:
             self._fail_listener(e)
 
-    def _register_listener(self):
+    def _register_listener(self) -> None:
         try:
             _ = self._channel.basic_consume(
                 queue=self._queue,
@@ -419,7 +419,7 @@ class Listener(AsyncEventLoopMixin):
         if not self._listener_future.done():
             self._listener_future.set_result(None)
 
-    def _fail_listener(self, exception: Exception):
+    def _fail_listener(self, exception: Exception) -> None:
         if not self._listener_future.done():
             self._listener_future.set_exception(exception)
 
@@ -432,7 +432,7 @@ class Listener(AsyncEventLoopMixin):
         method: Basic.Deliver,
         properties: BasicProperties,
         body: bytes,
-    ):
+    ) -> None:
         del channel
 
         if properties.headers is not None:
@@ -486,7 +486,7 @@ class Listener(AsyncEventLoopMixin):
 
     def _on_submitted_listener_error(
         self, listener_message_meta: ListenerMessageMeta, future: Future
-    ):
+    ) -> None:
         if future.cancelled():
             return
 
@@ -499,7 +499,9 @@ class Listener(AsyncEventLoopMixin):
                 message=f"error occured while submitting listener callback on listener `{listener_message_meta.listener_name}` and queue `{self._queue}`",
             )
 
-    def _parse_and_execute(self, listener_message_meta: ListenerMessageMeta):
+    def _parse_and_execute(
+        self, listener_message_meta: ListenerMessageMeta
+    ) -> None:
         try:
             listener_method_args, listener_method_kwargs = self._parse_args(
                 listener_message_meta
@@ -543,7 +545,7 @@ class Listener(AsyncEventLoopMixin):
         except:
             message = listener_message_meta.body.decode()
 
-        assigned_args = []
+        assigned_args: list[Any] = []
         listener_method_args = []
         listener_method_kwargs = {}
         next_positional_arg = 0
@@ -578,8 +580,8 @@ class Listener(AsyncEventLoopMixin):
                         )
 
                     if (
-                        parameter.kind is Parameter.POSITIONAL_ONLY
-                        or parameter.kind is Parameter.POSITIONAL_OR_KEYWORD
+                        parameter.kind == Parameter.POSITIONAL_ONLY
+                        or parameter.kind == Parameter.POSITIONAL_OR_KEYWORD
                     ):
                         if (
                             listener_context is not None
@@ -614,7 +616,7 @@ class Listener(AsyncEventLoopMixin):
                             raise ValueError(
                                 f"argument {parameter.name} has no default"
                             )
-                    elif parameter.kind is Parameter.VAR_POSITIONAL:
+                    elif parameter.kind == Parameter.VAR_POSITIONAL:
                         listener_method_args.extend(
                             [
                                 self._validate_parameter(
@@ -626,10 +628,7 @@ class Listener(AsyncEventLoopMixin):
                         )
 
                         next_positional_arg += len(args[next_positional_arg:])
-                    elif (
-                        parameter.kind is Parameter.KEYWORD_ONLY
-                        or parameter.kind is Parameter.POSITIONAL_OR_KEYWORD
-                    ):
+                    elif parameter.kind == Parameter.KEYWORD_ONLY:
                         if (
                             listener_context is not None
                             or dependency is not None
@@ -654,7 +653,7 @@ class Listener(AsyncEventLoopMixin):
                             raise ValueError(
                                 f"argument {parameter.name} has no default"
                             )
-                    elif parameter.kind is Parameter.VAR_KEYWORD:
+                    elif parameter.kind == Parameter.VAR_KEYWORD:
                         listener_method_kwargs.update(
                             {
                                 k: self._validate_parameter(
@@ -741,7 +740,7 @@ class Listener(AsyncEventLoopMixin):
         self,
         listener_message_meta: ListenerMessageMeta,
         task_or_future: Task | Future,
-    ):
+    ) -> None:
         if task_or_future.cancelled():
             return
 
@@ -829,7 +828,7 @@ class Listener(AsyncEventLoopMixin):
 
     def _observe_listener_time(
         self, listener_message_meta: ListenerMessageMeta
-    ):
+    ) -> None:
         self.LISTENER_PROCESSING_LATENCY.labels(
             queue=self._queue, listener_name=listener_message_meta.listener_name
         ).observe(listener_message_meta.listener_start_time - time())
@@ -876,13 +875,13 @@ class Listener(AsyncEventLoopMixin):
         exception: BaseException,
         listener_message_meta: ListenerMessageMeta,
         message: str | None = None,
-    ):
+    ) -> None:
         context_dispose_callback = None
         rpc_reply = None
 
         if listener_message_meta.properties.reply_to is not None:
 
-            def on_context_disposed(context: ListenerContext):
+            def on_context_disposed(context: ListenerContext) -> None:
                 assert listener_message_meta.properties.reply_to is not None
                 assert context.rpc_reply is not None
 
@@ -954,7 +953,7 @@ class Listener(AsyncEventLoopMixin):
         listener_message_meta: ListenerMessageMeta,
         retry_policy: RetryPolicy,
         times_rejected: int,
-    ):
+    ) -> None:
         message_redelivery_delay = retry_policy.next_delay(times_rejected)
 
         self._logger.debug(
@@ -977,7 +976,7 @@ class Listener(AsyncEventLoopMixin):
         self,
         listener_message_meta: ListenerMessageMeta,
         times_rejected: int,
-    ):
+    ) -> None:
         self.loop.create_task(self._channel_pool.get()).add_done_callback(
             partial(
                 self._on_redelivery_channel_found,
@@ -991,7 +990,7 @@ class Listener(AsyncEventLoopMixin):
         listener_message_meta: ListenerMessageMeta,
         times_rejected: int,
         task: Task[BaseChannel],
-    ):
+    ) -> None:
         if task.cancelled():
             return
 
@@ -1044,7 +1043,7 @@ class Listener(AsyncEventLoopMixin):
         self,
         listener_message_meta: ListenerMessageMeta,
         response: Any,
-    ):
+    ) -> None:
         assert listener_message_meta.properties.reply_to is not None
 
         reponse_properties = BasicProperties(content_type="application/json")
@@ -1108,7 +1107,7 @@ class Listener(AsyncEventLoopMixin):
         response_body: bytes,
         response_properties: BasicProperties,
         task: Task[BaseChannel],
-    ):
+    ) -> None:
         if task.cancelled():
             return
 
@@ -1159,7 +1158,7 @@ class Consumer(Listener):
         queue: str,
         prefetch_count: int = 250,
         retry_policy: RetryPolicy | None = None,
-    ):
+    ) -> None:
         super().__init__(
             queue=queue,
             listener_name_header_key="target",
@@ -1170,7 +1169,7 @@ class Consumer(Listener):
     def consume(
         self, target: str | None = None, retry_policy: RetryPolicy | None = None
     ) -> Callable[[Callable], Callable]:
-        def wrapper(consumer_method: Callable):
+        def wrapper(consumer_method: Callable) -> Callable:
             if not callable(consumer_method):
                 raise TypeError(
                     f"consumer method argument not a callable, got {type(consumer_method)}"
@@ -1221,7 +1220,7 @@ def consume(
 
 
 class RpcWorker(Listener):
-    def __init__(self, queue: str, prefetch_count: int = 250):
+    def __init__(self, queue: str, prefetch_count: int = 250) -> None:
         super().__init__(
             queue=queue,
             listener_name_header_key="procedure",
@@ -1232,7 +1231,7 @@ class RpcWorker(Listener):
     def execute(
         self, procedure: str | None = None
     ) -> Callable[[Callable], Callable]:
-        def wrapper(worker_method: Callable):
+        def wrapper(worker_method: Callable) -> Callable:
             if not callable(worker_method):
                 raise TypeError(
                     f"worker method argument not a callable, got {type(worker_method)}"

qena_shared_lib/rabbitmq/_pool.py

@@ -21,14 +21,14 @@ class ChannelPool(AsyncEventLoopMixin):
     async def fill(
         self,
         connection: AsyncioConnection,
-    ):
+    ) -> None:
         self._connection = connection
 
         async with self._pool_lock:
             for _ in range(self._initial_pool_size):
                 await self._add_new_channel()
 
-    async def drain(self):
+    async def drain(self) -> None:
         async with self._pool_lock:
             for channel_base in self._pool.values():
                 with channel_base as channel:

qena_shared_lib/rabbitmq/_publisher.py

@@ -40,13 +40,13 @@ class Publisher:
         )
         self._logger = LoggerProvider.default().get_logger("rabbitmq.publisher")
 
-    async def publish_as_arguments(self, *args, **kwargs):
+    async def publish_as_arguments(self, *args: Any, **kwargs: Any) -> None:
         await self._get_channel_and_publish({"args": args, "kwargs": kwargs})
 
-    async def publish(self, message: Any | None = None):
+    async def publish(self, message: Any | None = None) -> None:
         await self._get_channel_and_publish(message)
 
-    async def _get_channel_and_publish(self, message: Any):
+    async def _get_channel_and_publish(self, message: Any) -> None:
         if self._blocked_connection_check_callback():
             raise RabbitMQBlockedError(
                 "rabbitmq broker is not able to accept message right now"

qena_shared_lib/rabbitmq/_rpc_client.py

@@ -2,7 +2,7 @@ from asyncio import Future, Lock
 from functools import partial
 from importlib import import_module
 from time import time
-from typing import Any, Callable
+from typing import Any, Callable, Generic, TypeVar
 from uuid import uuid4
 
 from pika import BasicProperties
@@ -27,26 +27,26 @@ __all__ = ["RpcClient"]
 
 class ExitHandler:
     _exiting = False
-    _rpc_futures = []
+    _rpc_futures: list[Future] = []
     _original_exit_handler: Callable
 
     @classmethod
-    def is_exising(cls):
+    def is_exising(cls) -> bool:
         return cls._exiting
 
     @classmethod
-    def add_rpc_future(cls, rpc_future: Future):
+    def add_rpc_future(cls, rpc_future: Future) -> None:
         cls._rpc_futures.append(rpc_future)
 
     @classmethod
-    def remove_rpc_future(cls, rpc_future: Future):
+    def remove_rpc_future(cls, rpc_future: Future) -> None:
         try:
             cls._rpc_futures.remove(rpc_future)
         except:
             pass
 
     @classmethod
-    def cancel_futures(cls):
+    def cancel_futures(cls) -> None:
         cls._exiting = True
 
         for rpc_future in cls._rpc_futures:
@@ -54,7 +54,7 @@ class ExitHandler:
                 rpc_future.cancel()
 
     @staticmethod
-    def patch_exit_handler():
+    def patch_exit_handler() -> None:
         try:
             Server = import_module("uvicorn.server").Server
         except ModuleNotFoundError:
@@ -64,11 +64,11 @@ class ExitHandler:
         Server.handle_exit = ExitHandler.handle_exit
 
     @staticmethod
-    def notify_clients():
+    def notify_clients() -> None:
         ExitHandler.cancel_futures()
 
     @staticmethod
-    def handle_exit(*args, **kwargs):
+    def handle_exit(*args: Any, **kwargs: Any) -> None:
         ExitHandler.notify_clients()
         ExitHandler._original_exit_handler(*args, **kwargs)
 
@@ -76,7 +76,10 @@ class ExitHandler:
 ExitHandler.patch_exit_handler()
 
 
-class RpcClient(AsyncEventLoopMixin):
+R = TypeVar("R")
+
+
+class RpcClient(Generic[R], AsyncEventLoopMixin):
     SUCCEEDED_RPC_CALLS = Counter(
         name="succeeded_rpc_calls",
         documentation="RPC calls made",
@@ -101,7 +104,7 @@ class RpcClient(AsyncEventLoopMixin):
         exchange: str | None = None,
         procedure: str | None = None,
         headers: dict[str, str] | None = None,
-        return_type: type | None = None,
+        return_type: type[R] | None = None,
         timeout: float = 15,
     ):
         self._routing_key = routing_key
@@ -117,23 +120,23 @@ class RpcClient(AsyncEventLoopMixin):
         self._blocked_connection_check_callback = (
             blocked_connection_check_callback
         )
-        self._rpc_future = None
-        self._rpc_call_start_time = None
+        self._rpc_future: Future[R] | None = None
+        self._rpc_call_start_time: float | None = None
         self._rpc_call_lock = Lock()
         self._rpc_call_pending = False
         self._logger = LoggerProvider.default().get_logger(
             "rabbitmq.rpc_client"
         )
 
-    async def call_with_arguments(self, *args, **kwargs) -> Any:
+    async def call_with_arguments(self, *args: Any, **kwargs: Any) -> R:
         return await self._get_channel_and_call(
             ({"args": args, "kwargs": kwargs})
         )
 
-    async def call(self, message: Any | None = None) -> Any:
+    async def call(self, message: Any | None = None) -> R:
         return await self._get_channel_and_call(message)
 
-    async def _get_channel_and_call(self, message: Any) -> Any:
+    async def _get_channel_and_call(self, message: Any) -> R:
         if self._blocked_connection_check_callback():
             raise RabbitMQBlockedError(
                 "rabbitmq broker is not able to accept message right now"
@@ -162,7 +165,7 @@ class RpcClient(AsyncEventLoopMixin):
 
         return await self._rpc_future
 
-    def _on_queue_declared(self, message: Any, method: Method):
+    def _on_queue_declared(self, message: Any, method: Method) -> None:
         try:
             self._rpc_reply_consumer_tag = self._channel.basic_consume(
                 queue=method.method.queue,
@@ -194,7 +197,7 @@ class RpcClient(AsyncEventLoopMixin):
             return
 
         if self._timeout > 0:
-            _ = self.loop.call_later(
+            self._timeout_timer_handle = self.loop.call_later(
                 delay=self._timeout, callback=self._on_timeout
             )
 
@@ -205,7 +208,7 @@ class RpcClient(AsyncEventLoopMixin):
             self._procedure,
         )
 
-    def _on_timeout(self):
+    def _on_timeout(self) -> None:
         self._finalize_call(
             exception=RabbitMQRpcRequestTimeoutError(
                 f"rpc worker didn't responed in a timely manner within `{self._timeout}` seconds"
@@ -218,7 +221,7 @@ class RpcClient(AsyncEventLoopMixin):
         method: Basic.Deliver,
         properties: BasicProperties,
         body: bytes,
-    ):
+    ) -> None:
         del channel, method
 
         if properties.correlation_id != self._correlation_id:
@@ -260,9 +263,12 @@ class RpcClient(AsyncEventLoopMixin):
 
     def _finalize_call(
         self,
-        response: Any | None = None,
+        response: Any = None,
         exception: BaseException | None = None,
-    ):
+    ) -> None:
+        if not self._timeout_timer_handle.cancelled():
+            self._timeout_timer_handle.cancel()
+
         self._rpc_call_pending = False
         self._channel.basic_cancel(self._rpc_reply_consumer_tag)
         self._channel_base.release()

qena_shared_lib/scheduler.py

@@ -36,8 +36,8 @@ class ScheduledTask:
     cron_expression: str
     zone_info: ZoneInfo | None
 
-    def __post_init__(self):
-        self._next_run_in = None
+    def __post_init__(self) -> None:
+        self._next_run_in: int | None = None
         self._ran = False
         self._paramters = {}
 
@@ -56,7 +56,7 @@ class ScheduledTask:
         return self._next_run_in
 
     @next_run_in.setter
-    def next_run_in(self, value: int):
+    def next_run_in(self, value: int) -> None:
         self._next_run_in = value
 
     @property
@@ -64,7 +64,7 @@ class ScheduledTask:
         return self._ran
 
     @ran.setter
-    def ran(self, value: bool):
+    def ran(self, value: bool) -> None:
         self._ran = value
 
     @property
@@ -73,7 +73,7 @@ class ScheduledTask:
 
 
 class Scheduler:
-    def __init__(self):
+    def __init__(self) -> None:
         self._scheduled_tasks: list[ScheduledTask] = []
 
     def schedule(
@@ -90,7 +90,7 @@ class Scheduler:
 
     def add_task(
         self, task: Callable, cron_expression: str, timezone: str | None
-    ):
+    ) -> None:
         self._scheduled_tasks.append(
             ScheduledTask(
                 task=task,
@@ -188,15 +188,17 @@ class ScheduleManager(AsyncEventLoopMixin):
         self,
         logstash: BaseLogstashSender,
         container: Container | None = None,
-    ):
+    ) -> None:
         self._container = container or Container()
         self._logstash = logstash
         self._scheduled_tasks: list[ScheduledTask] = []
-        self._next_run_in = None
-        self._scheduler_task = None
+        self._next_run_in: int | None = None
+        self._scheduler_task: Task | None = None
         self._logger = LoggerProvider.default().get_logger("schedule_manager")
 
-    def include_scheduler(self, scheduler: Scheduler | type[SchedulerBase]):
+    def include_scheduler(
+        self, scheduler: Scheduler | type[SchedulerBase]
+    ) -> None:
         if isinstance(scheduler, Scheduler):
             self._scheduled_tasks.extend(scheduler.scheduled_tasks)
 
@@ -229,7 +231,7 @@ class ScheduleManager(AsyncEventLoopMixin):
     def scheduled_task_count(self) -> int:
         return len(self._scheduled_tasks)
 
-    def start(self):
+    def start(self) -> None:
         if not self._aquired_lock():
             return
 
@@ -251,20 +253,20 @@ class ScheduleManager(AsyncEventLoopMixin):
         self._scheduler_task.add_done_callback(self._on_scheduler_done)
         self.SCHEDULE_MANAGER_STATE.state("running")
 
-    def use_schedulers(self):
+    def use_schedulers(self) -> None:
         for scheduler in [
             scheduler.register_scheduled_tasks()
             for scheduler in self._container.resolve_all(SchedulerBase)
         ]:
             self._scheduled_tasks.extend(scheduler.scheduled_tasks)
 
-    def stop(self):
+    def stop(self) -> None:
         if self._scheduler_task is not None and not self._scheduler_task.done():
             self._scheduler_task.cancel()
 
         self.SCHEDULE_MANAGER_STATE.state("stopped")
 
-    def _on_scheduler_done(self, task: Task):
+    def _on_scheduler_done(self, task: Task) -> None:
         if task.cancelled():
             return
 
@@ -306,7 +308,7 @@ class ScheduleManager(AsyncEventLoopMixin):
 
         return True
 
-    async def _run_scheduler(self):
+    async def _run_scheduler(self) -> None:
         while True:
             self._calculate_next_schedule()
             self._logger.debug(
@@ -344,7 +346,7 @@ class ScheduleManager(AsyncEventLoopMixin):
     def _resolve_dependencies(
         self, scheduled_task: ScheduledTask
     ) -> dict[str, Any]:
-        args = {}
+        args: dict[str, Any] = {}
 
         if self._container is None:
             return args
@@ -354,7 +356,7 @@ class ScheduleManager(AsyncEventLoopMixin):
 
         return args
 
-    def _on_task_done(self, task_or_future: Task | Future):
+    def _on_task_done(self, task_or_future: Task | Future) -> None:
         if task_or_future.cancelled():
             return
 
@@ -366,7 +368,7 @@ class ScheduleManager(AsyncEventLoopMixin):
                 exception=exception,
             )
 
-    def _calculate_next_schedule(self):
+    def _calculate_next_schedule(self) -> None:
         prev_run_in = self._next_run_in or 0
         self._next_run_in = None