Fix functools.partial() typehint breakage when using future annotations

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:

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