diff --git a/hamilton/function_modifiers/adapters.py b/hamilton/function_modifiers/adapters.py index 11133822d..8a80ba50d 100644 --- a/hamilton/function_modifiers/adapters.py +++ b/hamilton/function_modifiers/adapters.py @@ -854,8 +854,8 @@ def save_json_data(data: pd.DataFrame, json_path: str = "data/my_saved_data.json def validate(self, fn: Callable): """Validates that the function output is a dict type.""" - return_annotation = inspect.signature(fn).return_annotation - if return_annotation is inspect.Signature.empty: + return_annotation = typing.get_type_hints(fn)["return"] + if return_annotation is None: raise InvalidDecoratorException( f"Function: {fn.__qualname__} must have a return annotation." ) diff --git a/tests/function_modifiers/test_adapters.py b/tests/function_modifiers/test_adapters.py index 06ba57b6b..4904510de 100644 --- a/tests/function_modifiers/test_adapters.py +++ b/tests/function_modifiers/test_adapters.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import dataclasses from collections import Counter from typing import Any, Collection, Dict, List, Tuple, Type