From e0e763170fb5cb6ed3e8f1065e26a912ea86ff22 Mon Sep 17 00:00:00 2001
From: Walter Oggioni <oggioni.walter@gmail.com>
Date: Fri, 1 Nov 2024 09:21:01 +0800
Subject: [PATCH] small fixes to type annotations

---
 src/pwo/__init__.py |  7 +++++--
 src/pwo/maybe.py    |  3 ++-
 src/pwo/private.py  | 12 ++++++------
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/pwo/__init__.py b/src/pwo/__init__.py
index 7eaad4a..6a4629d 100644
--- a/src/pwo/__init__.py
+++ b/src/pwo/__init__.py
@@ -6,10 +6,12 @@ from .private import (
     ExceptionHandlerOutcome,
     tmpdir,
     decorator_with_kwargs,
-    classproperty
+    classproperty,
+    AsyncQueueIterator
 )
 from .maybe import Maybe
 from .notification import TopicManager, Subscriber
+
 __all__ = [
     'format_filesize',
     'async_retry',
@@ -21,5 +23,6 @@ __all__ = [
     'decorator_with_kwargs',
     'classproperty',
     'TopicManager',
-    'Subscriber'
+    'Subscriber',
+    'AsyncQueueIterator'
 ]
diff --git a/src/pwo/maybe.py b/src/pwo/maybe.py
index f0f6ad2..ef3da79 100644
--- a/src/pwo/maybe.py
+++ b/src/pwo/maybe.py
@@ -53,6 +53,7 @@ class Maybe(Generic[T]):
 
     def __or__(self, alt: Maybe[T]) -> Maybe[T]:
         return self if self.is_present else alt
+
     def or_else(self, alt: T) -> T:
         return self.value if self.is_present else alt
 
@@ -65,7 +66,7 @@ class Maybe(Generic[T]):
         else:
             raise supplier()
 
-    def or_else_get(self, supplier: Callable[[], T]) -> Maybe[T]:
+    def or_else_get(self, supplier: Callable[[], Optional[T]]) -> Maybe[T]:
         return self if self.is_present else Maybe.of_nullable(supplier())
 
     def if_present(self, callback: Callable[[T], U]) -> None:
diff --git a/src/pwo/private.py b/src/pwo/private.py
index 210b772..71e57ab 100644
--- a/src/pwo/private.py
+++ b/src/pwo/private.py
@@ -6,7 +6,7 @@ from inspect import signature
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from time import sleep
-from typing import Callable
+from typing import Callable, AsyncIterator
 
 
 def decorator_with_kwargs(decorator: Callable) -> Callable:
@@ -207,16 +207,16 @@ def classproperty(func):
     return ClassPropertyDescriptor(func)
 
 
-class AsyncQueueIterator:
-    _queue: Queue
+class AsyncQueueIterator[T]:
+    _queue: Queue[T]
 
-    def __init__(self, queue: Queue):
+    def __init__(self, queue: Queue[T]):
         self._queue = queue
 
-    def __aiter__(self):
+    def __aiter__(self) -> AsyncIterator[T]:
         return self
 
-    async def __anext__(self):
+    async def __anext__(self) -> [T]:
         item = await self._queue.get()
         if item is None:
             raise StopAsyncIteration