fix mypy linting

Author: Andrei Bratu

src/humanloop/decorators/flow.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from functools import wraps
-from typing import Any, Callable, Optional, TypeVar
+from typing import Any, Awaitable, Callable, Coroutine, Optional, TypeVar
 from typing_extensions import ParamSpec
 
 from opentelemetry.trace import Span, Tracer
@@ -123,7 +123,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
                         return func_output  # type: ignore [return-value]
 
         @wraps(func)
-        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
+        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Awaitable[Optional[R]]:
             span: Span
             with set_decorator_context(
                 DecoratorContext(
@@ -157,7 +157,8 @@ async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
                         log_error: Optional[str]
                         log_output_message: Optional[ChatMessage]
                         try:
-                            func_output = await func(*args, **kwargs)
+                            # Polymorphic decorator does not recognize the function is a coroutine
+                            func_output = await func(*args, **kwargs)  # type: ignore [misc]
                             if (
                                 isinstance(func_output, dict)
                                 and len(func_output.keys()) == 2
@@ -206,7 +207,8 @@ async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
                 version=FlowDict(**flow_kernel),  # type: ignore
                 callable=async_wrapper,
             )
-            return async_wrapper
+            # Polymorphic decorator does not recognize the function is a coroutine
+            return async_wrapper  # type: ignore [return-value]
 
         # If the decorated function is a sync function, return the sync wrapper
         wrapper.file = File(  # type: ignore

src/humanloop/decorators/tool.py

@@ -8,7 +8,7 @@
 from dataclasses import dataclass
 from functools import wraps
 from inspect import Parameter
-from typing import Any, Callable, Literal, Mapping, Optional, Sequence, TypeVar, TypedDict, Union
+from typing import Any, Awaitable, Callable, Coroutine, Literal, Mapping, Optional, Sequence, TypeVar, TypedDict, Union
 from typing_extensions import ParamSpec
 
 from opentelemetry.trace import Tracer
@@ -114,7 +114,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
                 return func_output
 
         @wraps(func)
-        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
+        async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Awaitable[Optional[R]]:
             evaluation_context = get_evaluation_context()
             if evaluation_context is not None:
                 if evaluation_context.path == path:
@@ -135,7 +135,8 @@ async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
 
                 func_output: Optional[R]
                 try:
-                    func_output = await func(*args, **kwargs)
+                    # Polymorphic decorator does not recognize the function is a coroutine
+                    func_output = await func(*args, **kwargs)  # type: ignore [misc]
                     log_output = process_output(
                         func=func,
                         output=func_output,
@@ -168,7 +169,8 @@ async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
                 )
 
                 # Return the output of the decorated function
-                return func_output
+                # Polymorphic decorator does not recognize the function is a coroutine
+                return func_output  # type: ignore [return-value]
 
         # If the decorated function is an async function, return the async wrapper
         if asyncio.iscoroutinefunction(func):
@@ -178,7 +180,8 @@ async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> Optional[R]:
                 version=tool_kernel,
                 callable=async_wrapper,
             )
-            return async_wrapper
+            # Polymorphic decorator does not recognize the function is a coroutine
+            return async_wrapper  # type: ignore [return-value]
 
         # If the decorated function is a sync function, return the sync wrapper
         wrapper.file = File(  # type: ignore

tests/conftest.py

@@ -126,8 +126,12 @@ def hl_test_exporter() -> HumanloopSpanExporter:
     Test Exporter where HTTP calls to Humanloop API
     are mocked.
     """
-    client = MagicMock()
-    exporter = HumanloopSpanExporter(client=client)
+    hl_client_headers = MagicMock()
+    hl_client_base_url = MagicMock()
+    exporter = HumanloopSpanExporter(
+        hl_client_headers=hl_client_headers,
+        hl_client_base_url=hl_client_base_url,
+    )
     return exporter