Adds Typeddict Extract fields subclass type check and test for it

Author: skrawcz

hamilton/function_modifiers/expanders.py

@@ -8,7 +8,7 @@
 import typing_extensions
 import typing_inspect
 
-from hamilton import node, registry
+from hamilton import htypes, node, registry
 from hamilton.dev_utils import deprecation
 from hamilton.function_modifiers import base
 from hamilton.function_modifiers.dependencies import (
@@ -772,8 +772,15 @@ def validate(self, fn: Callable):
             else:
                 # check that fields is a subset of TypedDict that is defined
                 typed_dict_fields = typing.get_type_hints(output_type)
-                for k, v in self.fields.items():
-                    if typed_dict_fields.get(k, None) != v:
+                for field_name, field_type in self.fields.items():
+                    expected_type = typed_dict_fields.get(field_name, None)
+                    if expected_type == field_type:
+                        pass  # we're definitely good
+                    elif expected_type is not None and htypes.custom_subclass_check(
+                        field_type, expected_type
+                    ):
+                        pass
+                    else:
                         raise base.InvalidDecoratorException(
                             f"Error {self.fields} did not match a subset of the TypedDict annotation's fields {typed_dict_fields}."
                         )

tests/function_modifiers/test_expanders.py

@@ -340,6 +340,41 @@ def return_dict() -> return_type:
     annotation.validate(return_dict)
 
 
+class SomeObject:
+    pass
+
+
+class InheritedObject(SomeObject):
+    pass
+
+
+class MyDictInheritance(TypedDict):
+    test: SomeObject
+    test2: str
+
+
+class MyDictInheritanceBadCase(TypedDict):
+    test: InheritedObject
+    test2: str
+
+
+def test_extract_fields_validate_happy_inheritance():
+    def return_dict() -> MyDictInheritance:
+        return {}
+
+    annotation = function_modifiers.extract_fields({"test": InheritedObject})
+    annotation.validate(return_dict)
+
+
+def test_extract_fields_validate_not_subclass():
+    def return_dict() -> MyDictInheritanceBadCase:
+        return {}
+
+    annotation = function_modifiers.extract_fields({"test": SomeObject})
+    with pytest.raises(base.InvalidDecoratorException):
+        annotation.validate(return_dict)
+
+
 @pytest.mark.parametrize(
     "return_type",
     [(int), (list), (np.ndarray), (pd.DataFrame), (MyDictBad)],