update

Author: ArnavBalyan

hamilton/function_modifiers/adapters.py

@@ -872,13 +872,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).get("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):
+        origin = typing.get_origin(return_annotation)
+        if return_annotation is not dict and origin is not dict:
             raise InvalidDecoratorException(f"Function: {fn.__qualname__} must return a dict.")
 
     def generate_nodes(self, fn: Callable, config) -> list[node.Node]:

tests/function_modifiers/test_adapters.py

@@ -835,6 +835,17 @@ def test_dataloader_future_annotations():
     assert custom_subclass_check(fg["sample_dataloader"].type, list)
 
 
+def test_datasaver_future_annotations():
+    from tests.resources import nodes_with_future_annotation
+
+    fn_to_collect = nodes_with_future_annotation.sample_datasaver
+    fg = graph.create_function_graph(
+        ad_hoc_utils.create_temporary_module(fn_to_collect),
+        config={},
+    )
+    assert "sample_datasaver" in fg
+
+
 def test_datasaver():
     annotation = datasaver()
     (node1,) = annotation.generate_nodes(correct_ds_function, {})

tests/resources/nodes_with_future_annotation.py

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from hamilton.function_modifiers import dataloader
+from hamilton.function_modifiers import dataloader, datasaver
 from hamilton.htypes import Collect, Parallelizable
 
 """Tests future annotations with common node types"""
@@ -41,3 +41,9 @@ def collected(standard: Collect[int]) -> int:
 def sample_dataloader() -> tuple[list[str], dict]:
     """Grouping here as the rest test annotations"""
     return ["a", "b", "c"], {}
+
+
+@datasaver()
+def sample_datasaver(standard: int) -> dict:
+    """Grouping here as the rest test annotations"""
+    return {"saved": standard}