fix: improve component.output_types decorator type hinting to support run_async methods (#9102)

* improve output_types type hinting

* better name

* docstrings

Author: anakin87

haystack/core/component/component.py

@@ -70,13 +70,13 @@
 """
 
 import inspect
-from collections.abc import Callable
+from collections.abc import Callable, Coroutine
 from contextlib import contextmanager
 from contextvars import ContextVar
 from copy import deepcopy
 from dataclasses import dataclass
 from types import new_class
-from typing import Any, Dict, Optional, Protocol, Type, TypeVar, runtime_checkable
+from typing import Any, Dict, Optional, Protocol, Type, TypeVar, Union, runtime_checkable
 
 from typing_extensions import ParamSpec
 
@@ -88,8 +88,10 @@
 
 logger = logging.getLogger(__name__)
 
-P = ParamSpec("P")
-R = TypeVar("R", bound=Dict[str, Any])
+RunParamsT = ParamSpec("RunParamsT")
+SyncRunReturnT = TypeVar("SyncRunReturnT", bound=Dict[str, Any])
+AsyncRunReturnT = TypeVar("AsyncRunReturnT", bound=Coroutine[Any, Any, Dict[str, Any]])
+RunReturnT = Union[SyncRunReturnT, AsyncRunReturnT]
 
 
 @dataclass
@@ -447,12 +449,13 @@ def run(self, value: int):
             instance, {name: OutputSocket(name=name, type=type_) for name, type_ in types.items()}, OutputSocket
         )
 
-    def output_types(self, **types: Any) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    def output_types(
+        self, **types: Any
+    ) -> Callable[[Callable[RunParamsT, RunReturnT]], Callable[RunParamsT, RunReturnT]]:
         """
         Decorator factory that specifies the output types of a component.
 
         Use as:
-
         ```python
         @component
         class MyComponent:
@@ -462,7 +465,7 @@ def run(self, value: int):
         ```
         """
 
-        def output_types_decorator(run_method: Callable[P, R]) -> Callable[P, R]:
+        def output_types_decorator(run_method: Callable[RunParamsT, RunReturnT]) -> Callable[RunParamsT, RunReturnT]:
             """
             Decorator that sets the output types of the decorated method.
 

releasenotes/notes/output-types-decorator-better-type-hinting-51bb3f28f55954b4.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    Improved type hinting for the `component.output_types` decorator. The type hinting for the decorator was originally
+    introduced  to avoid overshadowing the type hinting of the `run` method and allow proper static type checking.
+    This update extends support to asynchronous `run_async` methods.