Typed dict extract fields

[skrawcz]

Adds typeddict support for extract_fields

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:

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.

Changes

• extract_fields decorator behavior -> it can now handle TypedDict

How I tested this

• locally
• unit tests

Notes

• fixes @extract_fields - support TypedDict as return type #1252

Checklist

• PR has an informative and human-readable title (this will be pulled into the release notes)
• Changes are limited to a single goal (no scope creep)
• Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
• Any change in functionality is tested
• New functions are documented (with a description, list of inputs, and expected output)
• Placeholder code is flagged / future TODOs are captured in comments
• Project documentation has been updated if adding/changing functionality.

[ellipsis-dev]

❌ Changes requested. Reviewed everything up to 1782dcf in 33 seconds

More details

• Looked at 130 lines of code in 2 files
• Skipped 0 files when reviewing.
• Skipped posting 2 drafted comments based on config settings.

1. tests/function_modifiers/test_expanders.py:365

• Draft comment:
  Consider adding a test case to check that an InvalidDecoratorException is raised when the specified fields do not match the TypedDict fields. This will ensure that the validation logic for TypedDicts is thoroughly tested.
• Reason this comment was not posted:
  Comment looked like it was already resolved.

2. hamilton/function_modifiers/expanders.py:777

• Draft comment:
  The logic for checking if fields is a subset of the TypedDict fields is repeated. Consider refactoring this into a separate function to adhere to the DRY principle.
• Reason this comment was not posted:
  Marked as duplicate.

Workflow ID: wflow_6h5DrbBG53z73ihx

---

Want Ellipsis to fix these issues? Tag @ellipsis-dev in a comment. You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

[ellipsis-dev]

The logic for checking if fields is a subset of the TypedDict fields is repeated. Consider refactoring this into a separate function to adhere to the DRY principle.

[ellipsis-dev]

👍 Looks good to me! Incremental review on f0bd568 in 17 seconds

More details

• Looked at 20 lines of code in 1 files
• Skipped 0 files when reviewing.
• Skipped posting 2 drafted comments based on config settings.

1. hamilton/function_modifiers/expanders.py:777

• Draft comment:
  The _validate_extract_fields function call has been moved to ensure it validates self.fields after it is potentially set from a TypedDict. This change ensures that the fields are validated correctly when using TypedDict.
• Reason this comment was not posted:
  Confidence changes required: 10%
  The PR modifies the validate method in extract_fields to handle TypedDict. The change moves the _validate_extract_fields call to ensure it is executed after setting self.fields when TypedDict is used. This ensures validation is performed on the correct fields. The change seems correct and aligns with the PR description.

2. hamilton/function_modifiers/expanders.py:783

• Draft comment:
  The error message here is misleading since TypedDict is also supported. Consider updating it to reflect that TypedDict is a valid return type.

                f"For extracting fields, output type must be a dict, typing.Dict, or TypedDict, not: {output_type}"

• Reason this comment was not posted:
  Comment was on unchanged code.

Workflow ID: wflow_K8SLBPmr15irZq3O

---

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.

[elijahbenizzy]

One minor point

[zilto]

Could we loosen the constraints to enable Mapping type?

[skrawcz]

Could we loosen the constraints to enable Mapping type?

if it's quick to do. Otherwise can punt for another PR.

[ellipsis-dev]

👍 Looks good to me! Incremental review on 4e8e239 in 30 seconds

More details

• Looked at 78 lines of code in 2 files
• Skipped 0 files when reviewing.
• Skipped posting 3 drafted comments based on config settings.

1. tests/function_modifiers/test_expanders.py:361

• Draft comment:
  Consider returning an instance of MyDictInheritance in return_dict to better simulate a real-world scenario and ensure the test covers both validation and extraction logic.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The test function test_extract_fields_validate_happy_inheritance is testing the validation of the extract_fields decorator with a TypedDict that involves inheritance. The test is correctly set up to ensure that the decorator can handle subclass relationships. However, the test function return_dict returns an empty dictionary, which might not be the best way to test the functionality. It would be more appropriate to return an instance of MyDictInheritance to better simulate a real-world scenario. This would ensure that the test is not only checking the validation logic but also the actual extraction logic.

2. tests/function_modifiers/test_expanders.py:369

• Draft comment:
  Consider returning an instance of MyDictInheritanceBadCase in return_dict to better simulate a real-world scenario and ensure the test covers both validation and extraction logic.
• Reason this comment was not posted:
  Confidence changes required: 50%
  The test function test_extract_fields_validate_not_subclass is designed to test the validation logic of the extract_fields decorator when the field type is not a subclass of the expected type. The test is correctly set up to ensure that the decorator raises an exception in this scenario. However, similar to the previous test, the function return_dict returns an empty dictionary. It would be more appropriate to return an instance of MyDictInheritanceBadCase to better simulate a real-world scenario and ensure the test covers both validation and extraction logic.

3. hamilton/function_modifiers/expanders.py:772

• Draft comment:
  The validate function in extract_fields is complex and would benefit from comments explaining the logic, especially around handling TypedDicts and field validation.
• Reason this comment was not posted:
  Comment was on unchanged code.

Workflow ID: wflow_qBjhjNI8059Htw65

---

You can customize Ellipsis with 👍 / 👎 feedback, review rules, user-specific overrides, quiet mode, and more.