This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 63d8acc795e [BEAM-19954] Resolve string annotations in callable type
hints (#40625)
63d8acc795e is described below
commit 63d8acc795e4ea00c88b28a892b5fa1d3b528fcb
Author: GAURAV THAKUR <[email protected]>
AuthorDate: Wed Oct 21 21:45:39 2026 +0530
[BEAM-19954] Resolve string annotations in callable type hints (#40625)
* [BEAM-19954] Resolve string annotations in callable type hints
Use typing.get_type_hints for annotated functions and class constructors so
forward references are converted to their intended Beam type hints. Preserve
the existing Any fallback when annotations cannot be resolved.
Add regression coverage for string and future annotations, methods,
constructors, and unresolved names.
* Apply Beam YAPF formatting to forward-reference fix
Match the formatter's expected layout for the return annotation and
postponed-annotation regression test.
* [BEAM-19954] Resolve forward references independently and correct ParDo
hint coverage
* Mark intentionally unresolved forward references in tests
---
sdks/python/apache_beam/typehints/decorators.py | 76 +++++++++
.../apache_beam/typehints/decorators_test.py | 189 +++++++++++++++++++++
.../apache_beam/typehints/typed_pipeline_test.py | 5 +-
3 files changed, 268 insertions(+), 2 deletions(-)
diff --git a/sdks/python/apache_beam/typehints/decorators.py
b/sdks/python/apache_beam/typehints/decorators.py
index bc44ec745d5..d3547f53568 100644
--- a/sdks/python/apache_beam/typehints/decorators.py
+++ b/sdks/python/apache_beam/typehints/decorators.py
@@ -80,11 +80,13 @@ defined, or before importing a module containing
type-hinted functions.
# pytype: skip-file
# ruff: noqa: UP006
+import functools
import inspect
import itertools
import logging
import traceback
import types
+import typing
from typing import Any
from typing import Callable
from typing import Dict
@@ -131,6 +133,67 @@ _ANY_VAR_KEYWORD = typehints.Dict[typehints.Any,
typehints.Any]
_disable_from_callable = False # pylint: disable=invalid-name
+def _contains_forward_ref(annotation):
+ if isinstance(annotation, (str, typing.ForwardRef)):
+ return True
+ origin = get_origin(annotation)
+ if origin is Literal:
+ # Literal strings are values, not type names.
+ return False
+ args = get_args(annotation)
+ if origin is typing.Annotated:
+ args = args[:1] # The remaining arguments are metadata.
+ return any(_contains_forward_ref(arg) for arg in args)
+
+
+def _annotation_namespaces(fn):
+ """Find the implementation and namespaces behind inspect.signature(fn)."""
+ target = fn
+ while isinstance(target, functools.partial):
+ target = target.func
+ localns = None
+ if inspect.isclass(target):
+ localns = dict(vars(target))
+ localns[target.__name__] = target
+ if target.__init__ is not object.__init__:
+ target = target.__init__
+ else:
+ target = target.__new__
+ elif inspect.ismethod(target):
+ owner = target.__self__
+ owner = owner if inspect.isclass(owner) else type(owner)
+ localns = dict(vars(owner))
+ localns[owner.__name__] = owner
+ elif not (inspect.isfunction(target) or inspect.isbuiltin(target)):
+ owner = type(target)
+ target = target.__call__
+ localns = dict(vars(owner))
+ localns[owner.__name__] = owner
+ target = inspect.unwrap(target)
+ if inspect.ismethod(target):
+ target = target.__func__
+ return getattr(target, '__globals__', {}), localns
+
+
+def _resolve_annotation(annotation, globalns, localns):
+ if not _contains_forward_ref(annotation):
+ return annotation
+
+ # get_type_hints resolves nested ForwardRefs and uses public typing APIs.
+ # A one-hint callable lets an unresolved name leave other hints intact.
+ def one_hint():
+ pass
+
+ one_hint.__annotations__ = {'value': annotation}
+ try:
+ return typing.get_type_hints(
+ one_hint, globalns=globalns, localns=localns,
+ include_extras=True)['value']
+ except (AttributeError, NameError, SyntaxError, TypeError):
+ # convert_to_beam_type maps unresolved strings and ForwardRefs to Any.
+ return annotation
+
+
def get_signature(func):
"""Like inspect.signature(), but supports Py2 as well.
@@ -334,6 +397,19 @@ class IOTypeHints(NamedTuple):
for param in signature.parameters.values()) and
signature.return_annotation == signature.empty):
return None
+ if (any(_contains_forward_ref(param.annotation)
+ for param in signature.parameters.values()) or
+ _contains_forward_ref(signature.return_annotation)):
+ globalns, localns = _annotation_namespaces(fn)
+ signature = signature.replace(
+ parameters=[
+ param.replace(
+ annotation=_resolve_annotation(
+ param.annotation, globalns, localns))
+ for param in signature.parameters.values()
+ ],
+ return_annotation=_resolve_annotation(
+ signature.return_annotation, globalns, localns))
input_args = []
input_kwargs = {}
for param in signature.parameters.values():
diff --git a/sdks/python/apache_beam/typehints/decorators_test.py
b/sdks/python/apache_beam/typehints/decorators_test.py
index 73efea9e9ab..3e9c3f18270 100644
--- a/sdks/python/apache_beam/typehints/decorators_test.py
+++ b/sdks/python/apache_beam/typehints/decorators_test.py
@@ -256,6 +256,195 @@ class IOTypeHintsTest(unittest.TestCase):
}))
self.assertEqual(th.output_types, ((Tuple[Any, ...], ), {}))
+ def test_from_callable_resolves_string_annotations(self):
+ def fn(value: 'int') -> 'typehints.List[str]':
+ return [str(value)]
+
+ th = decorators.IOTypeHints.from_callable(fn)
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.output_types, ((List[str], ), {}))
+
+ def test_from_callable_pardo_output_is_an_iterable_of_elements(self):
+ def do_fn(element: 'int') -> 'typehints.List[typehints.List[str]]':
+ return [[str(element)] * 2]
+
+ th = decorators.IOTypeHints.from_callable(do_fn)
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.strip_iterable().output_types, ((List[str], ), {}))
+
+ def test_from_callable_resolves_future_annotations(self):
+ namespace = {'typehints': typehints}
+ exec(
+ 'from __future__ import annotations\n'
+ 'def fn(value: int) -> typehints.List[str]:\n'
+ ' return [str(value)]\n',
+ namespace)
+
+ th = decorators.IOTypeHints.from_callable(namespace['fn'])
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.output_types, ((List[str], ), {}))
+
+ def test_from_callable_resolves_class_constructor_string_annotations(self):
+ class Class(object):
+ def __init__(self, value: 'int'):
+ self.value = value
+
+ th = decorators.IOTypeHints.from_callable(Class)
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.output_types, ((Class, ), {}))
+
+ def test_from_callable_resolves_method_string_annotations(self):
+ class Class(object):
+ def process(self, value: 'int') -> 'typehints.List[str]':
+ return [str(value)]
+
+ th = decorators.IOTypeHints.from_callable(Class().process)
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.output_types, ((List[str], ), {}))
+
+ def test_from_callable_unresolved_string_annotations_fall_back_to_any(self):
+ def fn(value: 'MissingType') -> 'int': # noqa: F821
+ return 1
+
+ th = decorators.IOTypeHints.from_callable(fn)
+ self.assertEqual(th.input_types, ((Any, ), {}))
+ self.assertEqual(th.output_types, ((int, ), {}))
+
+ def test_from_callable_resolves_other_hints_after_missing_name(self):
+ namespace = {}
+ exec(
+ 'from __future__ import annotations\n'
+ 'def fn(missing: MissingType, value: int, *, name: str) -> int:\n'
+ ' return value\n',
+ namespace)
+
+ th = decorators.IOTypeHints.from_callable(namespace['fn'])
+ self.assertEqual(th.input_types, ((Any, int), {'name': str}))
+ self.assertEqual(th.output_types, ((int, ), {}))
+
+ def test_from_callable_resolves_nested_forward_references(self):
+ def fn(
+ values: typing.List['int'],
+ pair: typing.Tuple['str', int],
+ builtin: list['int'],
+ nested: typing.Dict[str, typing.List[typing.ForwardRef('int')]],
+ resolved: typing.List[int]) -> typing.List['str']:
+ return [str(values[0])]
+
+ th = decorators.IOTypeHints.from_callable(fn)
+ self.assertEqual(
+ th.input_types,
+ ((
+ List[int],
+ Tuple[str, int],
+ List[int],
+ Dict[str, List[int]],
+ List[int]), {}))
+ self.assertEqual(th.output_types, ((List[str], ), {}))
+
+ def test_from_callable_unresolved_nested_reference_falls_back_to_any(self):
+ def fn(values: typing.List['MissingType']) -> 'int': # noqa: F821
+ return len(values)
+
+ th = decorators.IOTypeHints.from_callable(fn)
+ self.assertEqual(th.input_types, ((List[Any], ), {}))
+ self.assertEqual(th.output_types, ((int, ), {}))
+
+ def test_from_callable_partial_with_string_annotations(self):
+ def fn(bound: 'int', value: 'str', *, option: 'int') -> 'str':
+ return value
+
+ th = decorators.IOTypeHints.from_callable(functools.partial(fn, 1))
+ self.assertEqual(th.input_types, ((str, ), {
+ 'option': int
+ }))
+ self.assertEqual(th.output_types, ((str, ), {}))
+
+ def test_from_callable_bound_and_unbound_methods_with_strings(self):
+ class Processor:
+ def process(self, value: 'int', *, label: 'str') -> 'str':
+ return label
+
+ unbound = decorators.IOTypeHints.from_callable(Processor.process)
+ self.assertEqual(unbound.input_types, ((Any, int), {'label': str}))
+ self.assertEqual(unbound.output_types, ((str, ), {}))
+ bound = decorators.IOTypeHints.from_callable(Processor().process)
+ self.assertEqual(bound.input_types, ((int, ), {
+ 'label': str
+ }))
+ self.assertEqual(bound.output_types, ((str, ), {}))
+
+ def test_from_callable_varargs_and_kwargs_with_strings(self):
+ def fn(
+ value: 'int',
+ *args: 'typing.Tuple[str]',
+ required: 'str',
+ **kwargs: 'typing.Dict[str, int]') -> 'int':
+ return value
+
+ th = decorators.IOTypeHints.from_callable(fn)
+ self.assertEqual(
+ th.input_types, ((int, Tuple[str]), {
+ 'required': str, 'kwargs': Dict[str, int]
+ }))
+ self.assertEqual(th.output_types, ((int, ), {}))
+
+ def test_from_callable_callable_instance_with_string_annotations(self):
+ class CallableClass:
+ def __call__(self, value: 'int') -> 'str':
+ return str(value)
+
+ th = decorators.IOTypeHints.from_callable(CallableClass())
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.output_types, ((str, ), {}))
+
+ def test_from_callable_wrapped_function_uses_original_namespace(self):
+ namespace = {'Alias': str}
+ exec('def original(value: "Alias") -> "int":\n return 1\n', namespace)
+
+ @functools.wraps(namespace['original'])
+ def wrapper(*args, **kwargs):
+ return namespace['original'](*args, **kwargs)
+
+ th = decorators.IOTypeHints.from_callable(wrapper)
+ self.assertEqual(th.input_types, ((str, ), {}))
+ self.assertEqual(th.output_types, ((int, ), {}))
+
+ def test_from_callable_class_new_and_inherited_init(self):
+ class WithNew:
+ def __new__(cls, value: 'int'):
+ return super().__new__(cls)
+
+ class NoConstructor:
+ pass
+
+ th = decorators.IOTypeHints.from_callable(WithNew)
+ self.assertEqual(th.input_types, ((int, ), {}))
+ self.assertEqual(th.output_types, ((WithNew, ), {}))
+ th = decorators.IOTypeHints.from_callable(NoConstructor)
+ self.assertEqual(th.input_types, ((), {}))
+ self.assertEqual(th.output_types, ((NoConstructor, ), {}))
+
+ def test_from_callable_resolves_class_local_names(self):
+ class Node:
+ def __init__(self, previous: 'Node'):
+ self.previous = previous
+
+ def link(self, next_node: 'Node') -> 'Node':
+ return next_node
+
+ th = decorators.IOTypeHints.from_callable(Node)
+ self.assertEqual(th.input_types, ((Node, ), {}))
+ self.assertEqual(th.output_types, ((Node, ), {}))
+ th = decorators.IOTypeHints.from_callable(Node(None).link)
+ self.assertEqual(th.input_types, ((Node, ), {}))
+ self.assertEqual(th.output_types, ((Node, ), {}))
+
+ def test_from_callable_annotated_preserves_metadata_during_resolution(self):
+ annotation = typing.Annotated[typing.List['int'], 'metadata']
+ resolved = decorators._resolve_annotation(annotation, globals(), locals())
+ self.assertEqual(resolved, typing.Annotated[typing.List[int], 'metadata'])
+
def test_from_callable_partial(self):
def fn(a: int) -> int:
return a
diff --git a/sdks/python/apache_beam/typehints/typed_pipeline_test.py
b/sdks/python/apache_beam/typehints/typed_pipeline_test.py
index 29cb97dd787..9a5efd3dead 100644
--- a/sdks/python/apache_beam/typehints/typed_pipeline_test.py
+++ b/sdks/python/apache_beam/typehints/typed_pipeline_test.py
@@ -399,7 +399,7 @@ class MainInputTest(unittest.TestCase):
_ = [1, 2, 3] | beam.ParDo(my_do_fn, a=1)
def test_typed_callable_string_literals(self):
- def do_fn(element: 'int') -> 'typehints.List[str]':
+ def do_fn(element: 'int') -> 'typehints.List[typehints.List[str]]':
return [[str(element)] * 2]
result = [1, 2] | beam.ParDo(do_fn)
@@ -439,7 +439,8 @@ class MainInputTest(unittest.TestCase):
def test_typed_dofn_string_literals(self):
class MyDoFn(beam.DoFn):
- def process(self, element: 'int') -> 'typehints.List[str]':
+ def process(
+ self, element: 'int') -> 'typehints.List[typehints.List[str]]':
return [[str(element)] * 2]
result = [1, 2] | beam.ParDo(MyDoFn())
