Threaded is now class

Author: mosquito

aiomisc/iterator_wrapper.py

@@ -9,7 +9,7 @@
 from types import TracebackType
 from typing import (
     Any, AsyncIterator, Awaitable, Callable, Deque, Generator, NoReturn,
-    Optional, Type, TypeVar, Union,
+    Optional, Type, TypeVar, Union, Generic, ParamSpec,
 )
 from weakref import finalize
 
@@ -19,8 +19,9 @@
 
 T = TypeVar("T")
 R = TypeVar("R")
+P = ParamSpec("P")
 
-GenType = Generator[T, R, None]
+GenType = Generator[T, None, None]
 FuncType = Callable[[], GenType]
 
 
@@ -144,7 +145,7 @@ class IteratorWrapperStatistic(Statistic):
     enqueued: int
 
 
-class IteratorWrapper(AsyncIterator, EventLoopMixin):
+class IteratorWrapper(Generic[P, T], AsyncIterator, EventLoopMixin):
     __slots__ = (
         "__channel",
         "__close_event",
@@ -155,9 +156,11 @@ class IteratorWrapper(AsyncIterator, EventLoopMixin):
     ) + EventLoopMixin.__slots__
 
     def __init__(
-        self, gen_func: FuncType,
+        self,
+        gen_func: Callable[P, Generator[T, None, None]],
         loop: Optional[asyncio.AbstractEventLoop] = None,
-        max_size: int = 0, executor: Optional[Executor] = None,
+        max_size: int = 0,
+        executor: Optional[Executor] = None,
         statistic_name: Optional[str] = None,
     ):
 
@@ -227,11 +230,9 @@ async def wait_closed(self) -> None:
             await asyncio.gather(self.__gen_task, return_exceptions=True)
 
     def _run(self) -> Any:
-        return self.loop.run_in_executor(
-            self.executor, self._in_thread,
-        )
+        return self.loop.run_in_executor(self.executor, self._in_thread)
 
-    def __aiter__(self) -> AsyncIterator[Any]:
+    def __aiter__(self) -> AsyncIterator[T]:
         if not self.loop.is_running():
             raise RuntimeError("Event loop is not running")
 
@@ -242,7 +243,7 @@ def __aiter__(self) -> AsyncIterator[Any]:
             self.__gen_task = gen_task
         return IteratorProxy(self, self.close)
 
-    async def __anext__(self) -> Awaitable[T]:
+    async def __anext__(self) -> T:
         try:
             item, is_exc = await self.__channel.get()
         except ChannelClosed:
@@ -269,13 +270,13 @@ async def __aexit__(
         await self.close()
 
 
-class IteratorProxy(AsyncIterator):
+class IteratorProxy(Generic[T], AsyncIterator):
     def __init__(
-        self, iterator: AsyncIterator,
+        self, iterator: AsyncIterator[T],
         finalizer: Callable[[], Any],
     ):
         self.__iterator = iterator
         finalize(self, finalizer)
 
-    def __anext__(self) -> Awaitable[Any]:
+    def __anext__(self) -> Awaitable[T]:
         return self.__iterator.__anext__()

aiomisc/service/base.py

@@ -9,7 +9,7 @@
 
 
 T = TypeVar("T")
-CoroutineType = Union[Coroutine[Any, Any, T], Generator[Any, None, T]]
+CoroutineType = Union[Coroutine[Any, Any, T]]
 
 
 class ServiceMeta(ABCMeta):

aiomisc/thread_pool.py

@@ -2,6 +2,7 @@
 import contextvars
 import inspect
 import logging
+import os
 import threading
 import time
 import warnings
@@ -12,21 +13,24 @@
 from queue import SimpleQueue
 from types import MappingProxyType
 from typing import (
-    Any, Awaitable, Callable, Coroutine, Dict, FrozenSet, Optional, Set, Tuple,
-    TypeVar,
+    Any, Awaitable, Callable, Coroutine, Dict, FrozenSet, Generic,
+    Optional, Set, Tuple, TypeVar, Generator, overload, Union
 )
 
 from ._context_vars import EVENT_LOOP
 from .compat import ParamSpec
 from .counters import Statistic
 from .iterator_wrapper import IteratorWrapper
 
-
 P = ParamSpec("P")
 T = TypeVar("T")
 F = TypeVar("F", bound=Callable[..., Any])
 log = logging.getLogger(__name__)
 
+THREADED_ITERABLE_DEFAULT_MAX_SIZE = int(
+    os.getenv("THREADED_ITERABLE_DEFAULT_MAX_SIZE", 1024)
+)
+
 
 def context_partial(
     func: F, *args: Any,
@@ -327,6 +331,7 @@ async def lazy_wrapper() -> T:
             return await loop.run_in_executor(
                 executor, partial(func, *args, **kwargs),
             )
+
         return lazy_wrapper()
 
 
@@ -340,22 +345,54 @@ async def _awaiter(future: asyncio.Future) -> T:
         raise
 
 
+class Threaded(Generic[P, T]):
+    __slots__ = ("func",)
+
+    def __init__(self, func: Callable[P, T]) -> None:
+        if asyncio.iscoroutinefunction(func):
+            raise TypeError("Can not wrap coroutine")
+        if inspect.isgeneratorfunction(func):
+            raise TypeError("Can not wrap generator function")
+        self.func = func
+
+    def sync_call(self, *args: P.args, **kwargs: P.kwargs) -> T:
+        return self.func(*args, **kwargs)
+
+    def async_call(self, *args: P.args, **kwargs: P.kwargs) -> Awaitable[T]:
+        return run_in_executor(func=self.func, args=args, kwargs=kwargs)
+
+    def __repr__(self) -> str:
+        return f"<Threaded {self.func.__name__} at {id(self):#x}>"
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Awaitable[T]:
+        return self.async_call(*args, **kwargs)
+
+    def __get__(self, instance: Any, owner: Optional[type] = None) -> Any:
+        if instance is None:
+            return self
+        return partial(self.async_call, instance)
+
+
+@overload
+def threaded(func: Callable[P, T]) -> Threaded[P, T]: ...
+
+
+@overload
 def threaded(
-    func: Callable[P, T],
-) -> Callable[P, Awaitable[T]]:
-    if asyncio.iscoroutinefunction(func):
-        raise TypeError("Can not wrap coroutine")
+    func: Callable[P, Generator[T, None, None]]
+) -> Callable[P, IteratorWrapper[P, T]]: ...
 
-    if inspect.isgeneratorfunction(func):
-        return threaded_iterable(func)
 
-    @wraps(func)
-    def wrap(
-        *args: P.args, **kwargs: P.kwargs,
-    ) -> Awaitable[T]:
-        return run_in_executor(func=func, args=args, kwargs=kwargs)
+def threaded(
+    func: Callable[P, T] | Callable[P, Generator[T, None, None]]
+) -> Threaded[P, T] | Callable[P, IteratorWrapper[P, T]]:
+    if inspect.isgeneratorfunction(func):
+        return threaded_iterable(
+            func,
+            max_size=THREADED_ITERABLE_DEFAULT_MAX_SIZE
+        )
 
-    return wrap
+    return Threaded(func)   # type: ignore
 
 
 def run_in_new_thread(
@@ -390,67 +427,156 @@ def run_in_new_thread(
     return future
 
 
+class ThreadedSeparate(Threaded[P, T]):
+    """
+    A decorator to run a function in a separate thread.
+    It returns an `asyncio.Future` that can be awaited.
+    """
+
+    def __init__(self, func: Callable[P, T], detach: bool = True) -> None:
+        super().__init__(func)
+        self.detach = detach
+
+    def async_call(self, *args: P.args, **kwargs: P.kwargs) -> Awaitable[T]:
+        return run_in_new_thread(
+            self.func, args=args, kwargs=kwargs, detach=self.detach,
+        )
+
+
 def threaded_separate(
-    func: F,
+    func: Callable[P, T],
     detach: bool = True,
-) -> Callable[..., Awaitable[Any]]:
+) -> ThreadedSeparate[P, T]:
     if isinstance(func, bool):
         return partial(threaded_separate, detach=detach)
 
     if asyncio.iscoroutinefunction(func):
         raise TypeError("Can not wrap coroutine")
 
-    @wraps(func)
-    def wrap(*args: Any, **kwargs: Any) -> Any:
-        future = run_in_new_thread(
-            func, args=args, kwargs=kwargs, detach=detach,
+    return ThreadedSeparate(func, detach=detach)
+
+
+class ThreadedIterable(Generic[P, T]):
+    def __init__(
+        self,
+        func: Callable[P, Generator[T, None, None]],
+        max_size: int = 0
+    ) -> None:
+        self.func = func
+        self.max_size = max_size
+
+    def sync_call(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> Generator[T, None, None]:
+        return self.func(*args, **kwargs)
+
+    def async_call(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> IteratorWrapper[P, T]:
+        return self.create_wrapper(*args, **kwargs)
+
+    def create_wrapper(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> IteratorWrapper[P, T]:
+        return IteratorWrapper(
+            partial(self.func, *args, **kwargs),
+            max_size=self.max_size,
         )
-        return future
 
-    return wrap
+    def __call__(
+        self,
+        *args: P.args,
+        **kwargs: P.kwargs
+    ) -> IteratorWrapper[P, T]:
+        return self.async_call(*args, **kwargs)
 
+    def __get__(
+        self,
+        instance: Any,
+        owner: Optional[type] = None
+    ) -> Any:
+        if instance is None:
+            return self
+        return partial(self.async_call, instance)
 
+
+@overload
 def threaded_iterable(
-    func: Optional[F] = None,
+    func: Callable[P, Generator[T, None, None]],
+    *,
     max_size: int = 0,
-) -> Any:
-    if isinstance(func, int):
-        return partial(threaded_iterable, max_size=func)
-    if func is None:
-        return partial(threaded_iterable, max_size=max_size)
+) -> "ThreadedIterable[P, T]": ...
 
-    @wraps(func)
-    def wrap(*args: Any, **kwargs: Any) -> Any:
-        return IteratorWrapper(
-            partial(func, *args, **kwargs),
-            max_size=max_size,
-        )
 
-    return wrap
+@overload
+def threaded_iterable(
+    *,
+    max_size: int = 0,
+) -> Callable[
+    [Callable[P, Generator[T, None, None]]], ThreadedIterable[P, T]]: ...
 
 
+def threaded_iterable(
+    func: Optional[Callable[P, Generator[T, None, None]]] = None,
+    *,
+    max_size: int = 0,
+) -> Union[
+    ThreadedIterable[P, T],
+    Callable[[Callable[P, Generator[T, None, None]]],
+    ThreadedIterable[P, T]]
+]:
+    if func is None:
+        return lambda f: ThreadedIterable(f, max_size=max_size)
+
+    return ThreadedIterable(func, max_size=max_size)
+
 class IteratorWrapperSeparate(IteratorWrapper):
     def _run(self) -> Any:
         return run_in_new_thread(self._in_thread)
 
 
+class ThreadedIterableSeparate(ThreadedIterable[P, T]):
+    def create_wrapper(
+        self, *args: P.args, **kwargs: P.kwargs
+    ) -> IteratorWrapperSeparate:
+        return IteratorWrapperSeparate(
+            partial(self.func, *args, **kwargs),
+            max_size=self.max_size,
+        )
+
+
+@overload
 def threaded_iterable_separate(
-    func: Optional[F] = None,
+    func: Callable[P, Generator[T, None, None]],
+    *,
     max_size: int = 0,
-) -> Any:
-    if isinstance(func, int):
-        return partial(threaded_iterable_separate, max_size=func)
+) -> "ThreadedIterable[P, T]": ...
+
+
+@overload
+def threaded_iterable_separate(
+    *,
+    max_size: int = 0,
+) -> Callable[
+    [Callable[P, Generator[T, None, None]]],
+    ThreadedIterableSeparate[P, T]
+]: ...
+
+
+def threaded_iterable_separate(
+    func: Optional[Callable[P, Generator[T, None, None]]] = None,
+    *,
+    max_size: int = 0,
+) -> Union[
+    ThreadedIterable[P, T],
+    Callable[[Callable[P, Generator[T, None, None]]],
+    ThreadedIterableSeparate[P, T]]
+]:
     if func is None:
-        return partial(threaded_iterable_separate, max_size=max_size)
+        return lambda f: ThreadedIterableSeparate(f, max_size=max_size)
 
-    @wraps(func)
-    def wrap(*args: Any, **kwargs: Any) -> Any:
-        return IteratorWrapperSeparate(
-            partial(func, *args, **kwargs),
-            max_size=max_size,
-        )
+    return ThreadedIterableSeparate(func, max_size=max_size)
 
-    return wrap
 
 
 class CoroutineWaiter:
@@ -509,4 +635,5 @@ def sync_await(
 ) -> T:
     async def awaiter() -> T:
         return await func(*args, **kwargs)
+
     return wait_coroutine(awaiter())