fix type hints for branched tasks and unhinted functions

Author: RichardZhu2

src/pyper/_core/decorators.py

@@ -15,7 +15,8 @@
 
 _P = ParamSpec('P')
 _R = t.TypeVar('R')
-_ArgsKwargs: t.TypeAlias = t.Optional[t.Tuple[t.Tuple[t.Any], t.Dict[str, t.Any]]]
+_Default = t.TypeVar('T', bound=t.NoReturn)  # Matches to no type hints
+ArgsKwargs: t.TypeAlias = t.Optional[t.Tuple[t.Tuple[t.Any], t.Dict[str, t.Any]]]
 
 
 class task:
@@ -28,19 +29,57 @@ class task:
         workers (int): Defines the number of workers to run the task
         throttle (int): Limits the number of results the task is able to produce when all consumers are busy
         multiprocess (bool): Allows the task to be multiprocessed (cannot be `True` for async tasks)
-        bind (tuple[args, kwargs]): Additional args and kwargs to bind to the task when defining a pipeline
+        bind (tuple[tuple, dict]): Additional args and kwargs to bind to the task when defining a pipeline
 
     Returns:
-        Pipeline: A `Pipeline` instance consisting of one task.
+        A `Pipeline` instance consisting of one task.
 
-    Example:
-    ```python
-    def f(x: int):
-        return x + 1
+    Examples:
+        ```python
+        def spam(x: int):
+            return x + 1
+
+        p = task(spam)
+
+        def ham(x: int):
+            return [x, x + 1, x + 2]
 
-    p = task(f, workers=10, multiprocess=True)
-    ```
+        p = task(ham, branch=True, workers=10)
+
+        async def eggs(x: int):
+            yield x
+            yield x + 1
+            yield x + 2
+
+        p = task(eggs, branch=True, throttle=1)
+        ```
     """
+    @t.overload
+    def __new__(
+            cls,
+            func: t.Callable[_P, _Default],
+            /,
+            *,
+            branch: bool = False,
+            join: bool = False,
+            workers: int = 1,
+            throttle: int = 0,
+            multiprocess: bool = False,
+            bind: ArgsKwargs = None) -> Pipeline[_P, _Default]: ...
+    
+    @t.overload
+    def __new__(
+            cls,
+            func: None = None,
+            /,
+            *,
+            branch: t.Literal[True],
+            join: bool = False,
+            workers: int = 1,
+            throttle: int = 0,
+            multiprocess: bool = False,
+            bind: ArgsKwargs = None) -> t.Type[_branched_partial_task]: ...
+    
     @t.overload
     def __new__(
             cls,
@@ -52,20 +91,20 @@ def __new__(
             workers: int = 1,
             throttle: int = 0,
             multiprocess: bool = False,
-            bind: _ArgsKwargs = None) -> t.Type[task]: ...
+            bind: ArgsKwargs = None) -> t.Type[task]: ...
 
     @t.overload
     def __new__(
             cls,
             func: t.Callable[_P, t.Union[t.Awaitable[t.Iterable[_R]], t.AsyncGenerator[_R]]],
             /,
             *,
-            branch: True,
+            branch: t.Literal[True],
             join: bool = False,
             workers: int = 1,
             throttle: int = 0,
             multiprocess: bool = False,
-            bind: _ArgsKwargs = None) -> AsyncPipeline[_P, _R]: ...
+            bind: ArgsKwargs = None) -> AsyncPipeline[_P, _R]: ...
 
     @t.overload
     def __new__(
@@ -78,20 +117,20 @@ def __new__(
             workers: int = 1,
             throttle: int = 0,
             multiprocess: bool = False,
-            bind: _ArgsKwargs = None) -> AsyncPipeline[_P, _R]: ...
+            bind: ArgsKwargs = None) -> AsyncPipeline[_P, _R]: ...
 
     @t.overload
     def __new__(
             cls,
             func: t.Callable[_P, t.Iterable[_R]],
             /,
             *,
-            branch: True,
+            branch: t.Literal[True],
             join: bool = False,
             workers: int = 1,
             throttle: int = 0,
             multiprocess: bool = False,
-            bind: _ArgsKwargs = None) -> Pipeline[_P, _R]: ...
+            bind: ArgsKwargs = None) -> Pipeline[_P, _R]: ...
 
     @t.overload
     def __new__(
@@ -104,7 +143,7 @@ def __new__(
             workers: int = 1,
             throttle: int = 0,
             multiprocess: bool = False,
-            bind: _ArgsKwargs = None) -> Pipeline[_P, _R]: ...
+            bind: ArgsKwargs = None) -> Pipeline[_P, _R]: ...
 
     def __new__(
             cls,
@@ -116,25 +155,44 @@ def __new__(
             workers: int = 1,
             throttle: int = 0,
             multiprocess: bool = False,
-            bind: _ArgsKwargs = 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, branch=branch, join=join, workers=workers, throttle=throttle, multiprocess=multiprocess, bind=bind)
         return Pipeline([Task(func=func, branch=branch, join=join, workers=workers, throttle=throttle, multiprocess=multiprocess, bind=bind)])
 
     @staticmethod
-    def bind(*args, **kwargs) -> _ArgsKwargs:
+    def bind(*args, **kwargs) -> ArgsKwargs:
         """Bind additional `args` and `kwargs` to a task.
 
         Example:
-        ```python
-        def f(x: int, y: int):
-            return x + y
+            ```python
+            def f(x: int, y: int):
+                return x + y
 
-        p = task(f, bind=task.bind(y=1))
-        p(x=1)
-        ```
+            p = task(f, bind=task.bind(y=1))
+            p(x=1)
+            ```
         """
         if not args and not kwargs:
             return None
         return args, kwargs
+
+
+class _branched_partial_task:
+    @t.overload
+    def __new__(cls, func: t.Callable[_P, _Default]) -> Pipeline[_P, _Default]: ...
+
+    @t.overload
+    def __new__(
+            cls,
+            func: t.Callable[_P, t.Union[t.Awaitable[t.Iterable[_R]], t.AsyncGenerator[_R]]]) -> AsyncPipeline[_P, _R]: ...
+        
+    @t.overload
+    def __new__(cls, func: t.Callable[_P, t.Iterable[_R]]) -> Pipeline[_P, _R]: ...
+
+    @t.overload
+    def __new__(cls, func: t.Callable[_P, _R]) -> Pipeline[_P, t.Any]: ...
+
+    def __new__(cls):
+        raise NotImplementedError