feat: ✨ Overloads for decorators

Author: remimd

README.md

@@ -1,7 +1,7 @@
 # python-cq
 
 [![CI](https://github.com/100nm/python-cq/actions/workflows/ci.yml/badge.svg)](https://github.com/100nm/python-cq)
-[![PyPI - Version](https://img.shields.io/pypi/v/python-cq.svg?color=blue)](https://pypi.org/project/python-cq/)
+[![PyPI - Version](https://img.shields.io/pypi/v/python-cq.svg?color=blue)](https://pypi.org/project/python-cq)
 [![PyPI - Downloads](https://img.shields.io/pypi/dm/python-cq.svg?color=blue)](https://pypistats.org/packages/python-cq)
 [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
 

cq/_core/dispatcher/pipe.py

@@ -1,6 +1,6 @@
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass, field
-from typing import Any, Self
+from typing import Any, Self, overload
 
 from cq._core.dispatcher.base import BaseDispatcher, Dispatcher
 
@@ -24,6 +24,24 @@ def __init__(self, dispatcher: Dispatcher[Any, Any]) -> None:
         self.__dispatcher = dispatcher
         self.__steps = []
 
+    @overload
+    def step[T](
+        self,
+        wrapped: PipeConverter[T, Any],
+        /,
+        *,
+        dispatcher: Dispatcher[T, Any] | None = ...,
+    ) -> PipeConverter[T, Any]: ...
+
+    @overload
+    def step[T](
+        self,
+        wrapped: None = ...,
+        /,
+        *,
+        dispatcher: Dispatcher[T, Any] | None = ...,
+    ) -> Callable[[PipeConverter[T, Any]], PipeConverter[T, Any]]: ...
+
     def step[T](
         self,
         wrapped: PipeConverter[T, Any] | None = None,

cq/_core/handler.py

@@ -5,7 +5,7 @@
 from functools import partial
 from inspect import Parameter, getmro, isclass
 from inspect import signature as inspect_signature
-from typing import Any, Protocol, Self, runtime_checkable
+from typing import Any, Protocol, Self, overload, runtime_checkable
 
 import injection
 
@@ -89,6 +89,27 @@ class HandlerDecorator[I, O]:
     manager: HandlerManager[I, O]
     injection_module: injection.Module = field(default_factory=injection.mod)
 
+    @overload
+    def __call__(
+        self,
+        input_or_handler_type: type[I],
+        /,
+    ) -> Callable[[HandlerType[[I], O]], HandlerType[[I], O]]: ...
+
+    @overload
+    def __call__(
+        self,
+        input_or_handler_type: HandlerType[[I], O],
+        /,
+    ) -> HandlerType[[I], O]: ...
+
+    @overload
+    def __call__(
+        self,
+        input_or_handler_type: None = ...,
+        /,
+    ) -> Callable[[HandlerType[[I], O]], HandlerType[[I], O]]: ...
+
     def __call__(
         self,
         input_or_handler_type: type[I] | HandlerType[[I], O] | None = None,