[BEAM-19954] Resolve string annotations in callable type hints

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():

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

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())