diff --git a/sdks/python/apache_beam/typehints/decorators.py b/sdks/python/apache_beam/typehints/decorators.py index bc44ec745d5a..d3547f535681 100644 --- a/sdks/python/apache_beam/typehints/decorators.py +++ b/sdks/python/apache_beam/typehints/decorators.py @@ -80,11 +80,13 @@ def foo((a, b)): # 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 @@ def foo((a, b)): _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 @@ def from_callable(cls, fn: Callable) -> Optional['IOTypeHints']: 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 73efea9e9abf..3e9c3f182709 100644 --- a/sdks/python/apache_beam/typehints/decorators_test.py +++ b/sdks/python/apache_beam/typehints/decorators_test.py @@ -256,6 +256,195 @@ def fn( })) 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 c97a1977e9c3..b98cc751a373 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 @@ def process(self, element: int, **side_inputs: typehints.Dict[str, str]) \ _ = [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 @@ def fn(element: float): 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())