From 422b9a751f6cb53f47474e26daf63429c0245849 Mon Sep 17 00:00:00 2001 From: kreczko Date: Fri, 28 Mar 2025 15:54:18 +0000 Subject: [PATCH 1/2] fix(issue #1297): add annotations to tests and add str check to validation --- hamilton/function_modifiers/adapters.py | 2 +- tests/function_modifiers/test_adapters.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/hamilton/function_modifiers/adapters.py b/hamilton/function_modifiers/adapters.py index 11133822d..2d181ffcf 100644 --- a/hamilton/function_modifiers/adapters.py +++ b/hamilton/function_modifiers/adapters.py @@ -860,7 +860,7 @@ def validate(self, fn: Callable): f"Function: {fn.__qualname__} must have a return annotation." ) # check that the return type is a dict - if return_annotation not in (dict, Dict): + if return_annotation not in (dict, Dict, "dict", "dict_"): raise InvalidDecoratorException(f"Function: {fn.__qualname__} must return a dict.") def generate_nodes(self, fn: Callable, config) -> List[node.Node]: 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 From e4f07b9a035b214a149ac630350f2aad14f9e1bd Mon Sep 17 00:00:00 2001 From: kreczko Date: Fri, 28 Mar 2025 17:08:07 +0000 Subject: [PATCH 2/2] refactor(datasaver.validate): use typing.get_type_hints instead of inspect.signature to resolve return type --- hamilton/function_modifiers/adapters.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/hamilton/function_modifiers/adapters.py b/hamilton/function_modifiers/adapters.py index 2d181ffcf..8a80ba50d 100644 --- a/hamilton/function_modifiers/adapters.py +++ b/hamilton/function_modifiers/adapters.py @@ -854,13 +854,13 @@ 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." ) # check that the return type is a dict - if return_annotation not in (dict, Dict, "dict", "dict_"): + if return_annotation not in (dict, Dict): raise InvalidDecoratorException(f"Function: {fn.__qualname__} must return a dict.") def generate_nodes(self, fn: Callable, config) -> List[node.Node]: