@extract_fields - support TypedDict as return type #1252
Description
Is your feature request related to a problem? Please describe.
I would like to be able to return an instance of TypedDict from a node to make my IDE warn me if I forgot some keys, for example:
import dataclasses
from typing import TypedDict
import hamilton.function_modifiers
from .api.kittens import Kitten, KittensApi
@dataclasses.dataclass
class T_memo:
last_kitten_timestamp: float
T_kittens = list[Kitten]
class T_kittens_and_updated_memo(TypedDict):
kittens: T_kittens
updated_memo: T_memo
@hamilton.function_modifiers.extract_fields(dict(kittens=T_kittens, updated_memo=T_memo))
def kittens_and_updated_memo(KittensApi: KittensApi, memo: T_memo) -> T_kittens_and_updated_memo:
# TODO: fetch some kittens
kittens: list[Kitten] = []
return T_kittens_and_updated_memo(kittens=kittens, updated_memo=memo)What I currently get is:
hamilton.function_modifiers.base.InvalidDecoratorException: For extracting fields, output type must be a dict or typing.Dict, not: <class 'kittens_project.pipeline.T_kittens_and_updated_memo'>
Describe the solution you'd like
In the implementation of @extract_fields, treat TypedDict as a normal dict.
Describe alternatives you've considered
Write this code instead:
from typing import Any
# ...
@hamilton.function_modifiers.extract_fields(dict(kittens=T_kittens, updated_memo=T_memo))
def kittens_and_updated_memo(KittensApi: KittensApi, memo: T_memo) -> dict[str, Any]: # 1️⃣ declare return type as just dict
# TODO: fetch some kittens
kittens: list[Kitten] = []
# ❌ bad idea - mypy / pyright won't warn about missing `updated_memo`
#return dict(kittens=kittens)
# 2️⃣ transform the result into a `dict` instance
return dict(T_kittens_and_updated_memo(kittens=kittens, updated_memo=memo))