Skip to content

Comments

fix: resolve string annotations from from __future__ import annotations#4508

Open
cluster2600 wants to merge 1 commit intozenml-io:developfrom
cluster2600:fix/annotations-2477
Open

fix: resolve string annotations from from __future__ import annotations#4508
cluster2600 wants to merge 1 commit intozenml-io:developfrom
cluster2600:fix/annotations-2477

Conversation

@cluster2600
Copy link

Problem

Using from __future__ import annotations (PEP 563) in a module that defines ZenML steps makes Python store all type annotations as plain strings instead of evaluated type objects.

inspect.signature() preserves these strings verbatim, so any downstream code that calls key.__mro__ — such as MaterializerRegistry.__getitem__ — breaks at pipeline compilation time.

Reported in #2477. Originally the failure was a cryptic AttributeError: 'str' object has no attribute '__mro__'. A later commit improved this to a RuntimeError with a better message, but the underlying limitation remained: from __future__ import annotations was entirely unsupported in step definition modules.

Solution

Add get_resolved_type_hints(), a thin wrapper around typing.get_type_hints(include_extras=True):

  • include_extras=True preserves Annotated[…] metadata (artifact names, configs, etc.).
  • Uses inspect.unwrap() to reach the original function even when it is wrapped by the @step decorator.
  • Returns an empty dict on failure (e.g. unresolvable forward references) so callers can fall back gracefully.

The helper is called in two places where annotations are first consumed:

  1. parse_return_type_annotations() (steps/utils.py) — step output types.
  2. validate_entrypoint_function() (steps/entrypoint_function_utils.py) — step input parameter types.

When resolution fails, the string annotation is kept; MaterializerRegistry.__getitem__ will then raise a RuntimeError with an updated, more accurate message explaining the actual root cause (unresolvable forward reference) rather than telling users to remove from __future__ import annotations.

Changes

File What changed
src/zenml/steps/utils.py Add get_type_hints import; add get_resolved_type_hints() helper; resolve string return_annotation in parse_return_type_annotations()
src/zenml/steps/entrypoint_function_utils.py Import get_resolved_type_hints; resolve string input annotations in validate_entrypoint_function()
src/zenml/materializers/materializer_registry.py Update error message to reflect the new root cause
tests/unit/steps/test_utils.py Tests for string return annotation resolution and get_resolved_type_hints error handling
tests/unit/steps/test_entrypoint_function_utils.py Tests for string input/output annotation resolution via validate_entrypoint_function

Testing

After this fix, a step like:

from __future__ import annotations

from zenml import step

@step
def my_step(x: int) -> str:
    return str(x)

compiles and runs correctly without any changes to the step code.

Closes #2477

…ons'

Using 'from __future__ import annotations' (PEP 563) makes Python store
all type annotations as plain strings instead of evaluated type objects.
inspect.signature() preserves these strings, so any downstream code that
calls key.__mro__ (e.g. MaterializerRegistry.__getitem__) would hit an
AttributeError on older ZenML versions, or a RuntimeError on more recent
ones that added an explicit string-type guard.

This commit adds get_resolved_type_hints(), a thin wrapper around
typing.get_type_hints(include_extras=True) that resolves forward
references while preserving Annotated metadata.  It is called in two
places where annotations are first consumed:

* parse_return_type_annotations() in steps/utils.py – for step output
  types.
* validate_entrypoint_function() in steps/entrypoint_function_utils.py –
  for step input parameter types.

Resolution failures (e.g. unresolvable forward references due to circular
imports or TYPE_CHECKING-only imports) are swallowed silently; the
original string annotation is kept, which will still surface a clear
RuntimeError from MaterializerRegistry.__getitem__ with a message updated
to reflect the actual root cause.

Fixes zenml-io#2477

Co-authored-by: Maxime Grenu <maxime.grenu@gmail.com>
@schustmi schustmi self-assigned this Feb 20, 2026
@bcdurak
Copy link
Contributor

bcdurak commented Feb 20, 2026

@cluster2600 , thank you for your contribution. We'll take a look.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants