add type hint system

Author: RichardZhu2

.coveragerc

@@ -2,4 +2,5 @@
 exclude_also = 
     # pragma: no cover
     if TYPE_CHECKING:
+    if t.TYPE_CHECKING:
     raise NotImplementedError

src/pyper/_core/decorators.py

@@ -1,47 +1,111 @@
 from __future__ import annotations
 
 import functools
-from typing import Callable, Dict, Optional, overload, Tuple
+import typing as t
 
-from .pipeline import Pipeline
+from .pipeline import AsyncPipeline, Pipeline
 from .task import Task
 
 
+_P = t.ParamSpec('P')
+_R = t.TypeVar('R')
+_ArgsKwargs: t.TypeAlias = t.Optional[t.Tuple[t.Tuple[t.Any], t.Dict[str, t.Any]]]
+
+
 class task:
     """Decorator class to transform a function into a `Task` object, and then initialize a `Pipeline` with this task.
     A Pipeline initialized in this way consists of one Task, and can be piped into other Pipelines.
 
     The behaviour of each task within a Pipeline is determined by the parameters:
-    `join`: allows the function to take all previous results as input, instead of single results
-    `concurrency`: runs the functions with multiple (async or threaded) workers
-    `throttle`: limits the number of results the function is able to produce when all consumers are busy
+    * `join`: allows the function to take all previous results as input, instead of single results
+    * `concurrency`: runs the functions with multiple (async or threaded) workers
+    * `throttle`: limits the number of results the function is able to produce when all consumers are busy
+    * `daemon`: determines whether threaded workers are daemon threads (cannot be True for async tasks)
+    * `bind`: additional args and kwargs to bind to the function when defining a pipeline
     """
-    @overload
-    def __new__(cls, func: None = None, /, *, join: bool = False, concurrency: int = 1, throttle: int = 0, daemon: bool = False, bind: Optional[Tuple[Tuple, Dict]] = None) -> Callable[..., Pipeline]:
-        """Enable type hints for functions decorated with `@task()`."""
+    @t.overload
+    def __new__(
+        cls,
+        func: None = None,
+        /,
+        *,
+        join: bool = False,
+        concurrency: int = 1,
+        throttle: int = 0,
+        daemon: bool = False,
+        bind: _ArgsKwargs = None
+    ) -> t.Type[task]: ...
 
-    @overload
-    def __new__(cls, func: Callable, /, *, join: bool = False, concurrency: int = 1, throttle: int = 0, daemon: bool = False, bind: Optional[Tuple[Tuple, Dict]] = None) -> Pipeline:
-        """Enable type hints for functions decorated with `@task`."""
+    @t.overload
+    def __new__(
+        cls,
+        func: t.Callable[_P, t.Awaitable[_R]],
+        /,
+        *,
+        join: bool = False,
+        concurrency: int = 1,
+        throttle: int = 0,
+        daemon: bool = False,
+        bind: _ArgsKwargs = None
+    ) -> AsyncPipeline[_P, _R]: ...
 
+    @t.overload
+    def __new__(
+        cls,
+        func: t.Callable[_P, t.AsyncGenerator[_R]],
+        /,
+        *,
+        join: bool = False,
+        concurrency: int = 1,
+        throttle: int = 0,
+        daemon: bool = False,
+        bind: _ArgsKwargs = None
+    ) -> AsyncPipeline[_P, _R]: ...
+    
+    @t.overload
+    def __new__(
+        cls,
+        func: t.Callable[_P, t.Generator[_R]],
+        /,
+        *,
+        join: bool = False,
+        concurrency: int = 1,
+        throttle: int = 0,
+        daemon: bool = False,
+        bind: _ArgsKwargs = None
+    ) -> Pipeline[_P, _R]: ...
+    
+    @t.overload
+    def __new__(
+        cls,
+        func: t.Callable[_P, _R],
+        /,
+        *,
+        join: bool = False,
+        concurrency: int = 1,
+        throttle: int = 0,
+        daemon: bool = False,
+        bind: _ArgsKwargs = None
+    ) -> Pipeline[_P, _R]: ...
+
     def __new__(
         cls,
-        func: Optional[Callable] = None,
+        func: t.Optional[t.Callable] = None,
         /,
         *,
         join: bool = False,
         concurrency: int = 1,
         throttle: int = 0,
         daemon: bool = False,
-        bind: Optional[Tuple[Tuple, Dict]] = None
+        bind: _ArgsKwargs = None
     ):
         # Classic decorator trick: @task() means func is None, @task without parentheses means func is passed. 
         if func is None:
             return functools.partial(cls, join=join, concurrency=concurrency, throttle=throttle, daemon=daemon, bind=bind)
         return Pipeline([Task(func=func, join=join, concurrency=concurrency, throttle=throttle, daemon=daemon, bind=bind)])
 
     @staticmethod
-    def bind(*args, **kwargs) -> Optional[Tuple[Tuple, Dict]]:
+    def bind(*args, **kwargs) -> _ArgsKwargs:
         """Utility method, to be used with `functools.partial`."""
         if not args and not kwargs:
             return None

src/pyper/_core/pipeline.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
 import inspect
-from typing import Callable, List, TYPE_CHECKING
+import typing as t
 
 from .async_helper.output import AsyncPipelineOutput
 from .sync_helper.output import PipelineOutput
 
-if TYPE_CHECKING:
+if t.TYPE_CHECKING:
     from .task import Task
 
 
-class Pipeline:
+_P = t.ParamSpec('P')
+_R = t.TypeVar('R')
+_P_Other = t.ParamSpec("P_Other")
+_R_Other = t.TypeVar("R_Other")
+
+
+class Pipeline(t.Generic[_P, _R]):
     """A sequence of at least 1 Tasks.
     
     Two pipelines can be piped into another via:
@@ -21,59 +27,83 @@ class Pipeline:
     ```
     """
 
-    def __new__(cls, tasks: List[Task]):
+    def __new__(cls, tasks: t.List[Task]):
         if any(task.is_async for task in tasks):
             instance = object.__new__(AsyncPipeline)
         else:
             instance = object.__new__(cls)
         instance.__init__(tasks=tasks)
         return instance
 
-    def __init__(self, tasks: List[Task]):
+    def __init__(self, tasks: t.List[Task]):
         self.tasks = tasks
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> t.Generator[_R]:
         """Return the pipeline output."""
         output = PipelineOutput(self)
         return output(*args, **kwargs)
+
+    @t.overload
+    def pipe(self: AsyncPipeline[_P, _R], other: AsyncPipeline[_P_Other, _R_Other]) -> AsyncPipeline[_P, _R_Other]: ...
+    
+    @t.overload
+    def pipe(self: AsyncPipeline[_P, _R], other: Pipeline[_P_Other, _R_Other]) -> AsyncPipeline[_P, _R_Other]: ...
 
-    def pipe(self, other) -> Pipeline:
+    @t.overload
+    def pipe(self, other: AsyncPipeline[_P_Other, _R_Other]) -> AsyncPipeline[_P, _R_Other]: ...
+    
+    @t.overload
+    def pipe(self, other: Pipeline[_P_Other, _R_Other]) -> Pipeline[_P, _R_Other]: ...
+    
+    def pipe(self, other: Pipeline):
         """Connect two pipelines, returning a new Pipeline."""
         if not isinstance(other, Pipeline):
             raise TypeError(f"{other} of type {type(other)} cannot be piped into a Pipeline")
         return Pipeline(self.tasks + other.tasks)
 
-    def __or__(self, other: Pipeline) -> Pipeline:
-        """Allow the syntax `pipeline1 | pipeline2`."""
+    @t.overload
+    def __or__(self: AsyncPipeline[_P, _R], other: AsyncPipeline[_P_Other, _R_Other]) -> AsyncPipeline[_P, _R_Other]: ...
+    
+    @t.overload
+    def __or__(self: AsyncPipeline[_P, _R], other: Pipeline[_P_Other, _R_Other]) -> AsyncPipeline[_P, _R_Other]: ...
+    
+    @t.overload
+    def __or__(self, other: AsyncPipeline[_P_Other, _R_Other]) -> AsyncPipeline[_P, _R_Other]: ...
+    
+    @t.overload
+    def __or__(self, other: Pipeline[_P_Other, _R_Other]) -> Pipeline[_P, _R_Other]: ...
+
+    def __or__(self, other: Pipeline):
+        """Connect two pipelines, returning a new Pipeline."""
         return self.pipe(other)
 
-    def consume(self, other: Callable) -> Callable:
+    def consume(self, other: t.Callable[..., _R_Other]) -> t.Callable[_P, _R_Other]:
         """Connect the pipeline to a consumer function (a callable that takes the pipeline output as input)."""
         if callable(other):
-            def consumer(*args, **kwargs):
+            def consumer(*args: _P.args, **kwargs: _P.kwargs) -> _R_Other:
                 return other(self(*args, **kwargs))
             return consumer
         raise TypeError(f"{other} must be a callable that takes a generator")
 
-    def __gt__(self, other: Callable) -> Callable:
-        """Allow the syntax `pipeline > consumer`."""
+    def __gt__(self, other: t.Callable[..., _R_Other]) -> t.Callable[_P, _R_Other]:
+        """Connect the pipeline to a consumer function (a callable that takes the pipeline output as input)."""
         return self.consume(other)
 
     def __repr__(self):
         return f"{self.__class__.__name__} {[task.func for task in self.tasks]}"
 
 
-class AsyncPipeline(Pipeline):
-    def __call__(self, *args, **kwargs):
+class AsyncPipeline(Pipeline[_P, _R]):
+    def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> t.AsyncGenerator[_R]:
         """Return the pipeline output."""
         output = AsyncPipelineOutput(self)
         return output(*args, **kwargs)
-        
-    def consume(self, other: Callable) -> Callable:
+    
+    def consume(self, other: t.Callable[..., _R_Other]) -> t.Callable[_P, _R_Other]:
         """Connect the pipeline to a consumer function (a callable that takes the pipeline output as input)."""
         if callable(other) and \
             (inspect.iscoroutinefunction(other) or inspect.iscoroutinefunction(other.__call__)):
-            async def consumer(*args, **kwargs):
+            async def consumer(*args: _P.args, **kwargs: _P.kwargs) -> _R_Other:
                 return await other(self(*args, **kwargs))
             return consumer
         raise TypeError(f"{other} must be an async callable that takes an async generator")