diff --git a/sdks/python/apache_beam/typehints/decorators.py b/sdks/python/apache_beam/typehints/decorators.py index d3547f535681..2eab07a4299d 100644 --- a/sdks/python/apache_beam/typehints/decorators.py +++ b/sdks/python/apache_beam/typehints/decorators.py @@ -854,6 +854,7 @@ def getcallargs_forhints(func, *type_args, **type_kwargs): # Might be raised due to too few or too many arguments. raise TypeCheckError(e) bound_args = bindings.arguments + namespaces = None for param in signature.parameters.values(): if param.name in bound_args: # Bound: unpack/convert variadic arguments. @@ -866,7 +867,16 @@ def getcallargs_forhints(func, *type_args, **type_kwargs): else: # Unbound: must have a default or be variadic. if param.annotation != param.empty: - bound_args[param.name] = param.annotation + annotation = param.annotation + if _contains_forward_ref(annotation): + # Resolve string annotations (e.g. from functools.partial-bound + # keywords under `from __future__ import annotations`) the same way + # IOTypeHints.from_callable does, so they compare equal to the + # resolved hints. + if namespaces is None: + namespaces = _annotation_namespaces(func) + annotation = _resolve_annotation(annotation, *namespaces) + bound_args[param.name] = convert_to_beam_type(annotation) elif param.kind == param.VAR_POSITIONAL: bound_args[param.name] = _ANY_VAR_POSITIONAL elif param.kind == param.VAR_KEYWORD: diff --git a/sdks/python/apache_beam/typehints/typed_pipeline_test.py b/sdks/python/apache_beam/typehints/typed_pipeline_test.py index 9a5efd3dead7..87e256c24dc7 100644 --- a/sdks/python/apache_beam/typehints/typed_pipeline_test.py +++ b/sdks/python/apache_beam/typehints/typed_pipeline_test.py @@ -20,6 +20,8 @@ # pytype: skip-file # ruff: noqa: UP006 +import functools +import types import typing import unittest from typing import Tuple @@ -405,6 +407,28 @@ def do_fn(element: 'int') -> 'typehints.List[typehints.List[str]]': result = [1, 2] | beam.ParDo(do_fn) self.assertEqual([['1', '1'], ['2', '2']], sorted(result)) + def test_typed_map_partial_with_future_annotations(self): + # Keywords bound by functools.partial are typed from their (string) + # annotations, which must resolve to the same types as the expected hints, + # including names referenced through an aliased module import. + class Config(object): + def __init__(self, scale): + self.scale = scale + + config_lib = types.ModuleType('config_lib') + config_lib.Config = Config + namespace = {'config_lib': config_lib} + exec( + 'from __future__ import annotations\n' + 'def scale(\n' + ' element: int, config: config_lib.Config, offset: int) -> int:\n' + ' return element * config.scale + offset\n', + namespace) + + fn = functools.partial(namespace['scale'], config=Config(2), offset=1) + result = [1, 2, 3] | beam.Map(fn) + self.assertEqual([3, 5, 7], sorted(result)) + def test_typed_ptransform_fn(self): # Test that type hints are propagated to the created PTransform. @beam.ptransform_fn