Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

@extract_fields - support TypedDict as return type #1252

Closed
earshinov opened this issue Dec 7, 2024 · 3 comments · Fixed by #1253
Closed

@extract_fields - support TypedDict as return type #1252

earshinov opened this issue Dec 7, 2024 · 3 comments · Fixed by #1253

Comments

@earshinov
Copy link

earshinov commented Dec 7, 2024 •
edited
Loading

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))
skrawcz added a commit that referenced this issue Dec 7, 2024
@skrawcz
Adds sketch of improving extract_fields with typeddict
413c282 
This in response to #1252.

We should be able to handle typeddict better.

This sketches some ideas:

1. field validation should happen in .validate() not the constructor.
2. extract_fields shouldn't need fields if the typeddict is the annotation type.
3. we properly check that typeddict can be a return type.
@skrawcz skrawcz linked a pull request Dec 7, 2024 that will close this issue
Typed dict extract fields #1253
Merged
7 tasks
@skrawcz
Copy link
Collaborator

skrawcz commented Dec 7, 2024

@earshinov thanks for the issue. I think we should be able to simplify things further and just enable:

@hamilton.function_modifiers.extract_fields()
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)

Thoughts?

I started a sketch in #1253 . Should be able to get it out this coming week some time.

@earshinov
Copy link
Author

I think we should be able to simplify things further and just enable

@skrawcz , yes, that would be wonderful!

skrawcz added a commit that referenced this issue Dec 8, 2024
@skrawcz
Adds sketch of improving extract_fields with typeddict
594f4e9 
This in response to #1252.

We should be able to handle typeddict better.

This sketches some ideas:

1. field validation should happen in .validate() not the constructor.
2. extract_fields shouldn't need fields if the typeddict is the annotation type.
3. we properly check that typeddict can be a return type.
@skrawcz skrawcz mentioned this issue Dec 8, 2024
Merged
7 tasks
skrawcz added a commit that referenced this issue Dec 12, 2024
@skrawcz
Adds sketch of improving extract_fields with typeddict
00e944d 
This in response to #1252.

We should be able to handle typeddict better.

This sketches some ideas:

1. field validation should happen in .validate() not the constructor.
2. extract_fields shouldn't need fields if the typeddict is the annotation type.
3. we properly check that typeddict can be a return type.
skrawcz added a commit that referenced this issue Dec 12, 2024
@skrawcz
Adds sketch of improving extract_fields with typeddict
15051e3 
This in response to #1252.

We should be able to handle typeddict better.

This sketches some ideas:

1. field validation should happen in .validate() not the constructor.
2. extract_fields shouldn't need fields if the typeddict is the annotation type.
3. we properly check that typeddict can be a return type.
skrawcz added a commit that referenced this issue Dec 12, 2024
@skrawcz
Adds sketch of improving extract_fields with typeddict
b982bf7 
This in response to #1252.

We should be able to handle typeddict better.

This sketches some ideas:

1. field validation should happen in .validate() not the constructor.
2. extract_fields shouldn't need fields if the typeddict is the annotation type.
3. we properly check that typeddict can be a return type.
skrawcz added a commit that referenced this issue Dec 12, 2024
@skrawcz
Adds typed dict extract fields support (#1253)
622866a 
Adds typeddict support for extract_fields

This enables more ergonomic usage of TypedDict with extract fields.

I skipped adding support for returning any `Mapping`. Though that should be an easy addition.

```python
from typing import TypedDict
from hamilton.function_modifiers import extract_fields

class MyDict(TypedDict):
    foo: str
    bar: int

@extract_fields()
def some_function() -> MyDict:
    return MyDict(foo="s", bar=1)
```
The above will automatically extract the fields foo and bar.

You can also do:

```python
from typing import TypedDict
from hamilton.function_modifiers import extract_fields

class MyDict(TypedDict):
    foo: str
    bar: int

@extract_fields({"foo": str})
def some_function()->MyDict:
    return MyDict(foo="s", bar=1)
```
To only expose a subset of the fields.


Squashed commits:

* Adds sketch of improving extract_fields with typeddict

This in response to #1252.

We should be able to handle typeddict better.

This sketches some ideas:

1. field validation should happen in .validate() not the constructor.
2. extract_fields shouldn't need fields if the typeddict is the annotation type.
3. we properly check that typeddict can be a return type.

* Adds typeddict tests

* Adding validation to cover all extract_fields paths

* Adds Typeddict Extract fields subclass type check and test for it
@skrawcz
Copy link
Collaborator

skrawcz commented Dec 12, 2024 •
edited
Loading

@earshinov published in sf-hamilton==1.85.0. Please verify and reopen if there's an issue.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

Typed dict extract fields DAGWorks-Inc/hamilton
2 participants
@earshinov @skrawcz