Restore previous handling of Callables, fix bare collections types (apache#34645)

jrmccluskey · web-flow · commit a847f066916e · 2025-04-17T15:02:07.000-04:00
* Restore previous handling of Callables

* fix bare type handling

* add parameterized tests
diff --git a/sdks/python/apache_beam/typehints/native_type_compatibility.py b/sdks/python/apache_beam/typehints/native_type_compatibility.py
@@ -364,7 +364,7 @@ def convert_to_beam_type(typ):
     # Only translate primitives and types from collections.abc and typing.
     return typ
   if (typ_module == 'collections.abc' and
-      typ.__origin__ not in _CONVERTED_COLLECTIONS):
+      getattr(typ, '__origin__', typ) not in _CONVERTED_COLLECTIONS):
     # TODO(https://github.com/apache/beam/issues/29135):
     # Support more collections types
     return typ
diff --git a/sdks/python/apache_beam/typehints/native_type_compatibility_test.py b/sdks/python/apache_beam/typehints/native_type_compatibility_test.py
@@ -263,6 +263,16 @@ def test_convert_to_beam_type_with_collections_types(self):
             collections.defaultdict[str, int],
             typehints.Dict[str, int]),
         ('count', collections.Counter[str, int], typehints.Dict[str, int]),
+        (
+            'bare callable',
+            collections.abc.Callable,
+            collections.abc.Callable,
+        ),
+        (
+            'parameterized callable',
+            collections.abc.Callable[[str], int],
+            collections.abc.Callable[[str], int],
+        ),
     ]
 
     for test_case in test_cases:
@@ -369,6 +379,50 @@ def test_convert_bare_types(self):
       converted_beam_type = convert_to_beam_type(typing_type)
       self.assertEqual(expected_beam_type, converted_beam_type, description)
 
+  def test_convert_bare_collections_types(self):
+    # Conversions for unsubscripted types that have implicit subscripts.
+    test_cases = [
+        ('bare list', list, typehints.List[typehints.TypeVariable('T')]),
+        (
+            'bare dict',
+            dict,
+            typehints.Dict[typehints.TypeVariable('KT'),
+                           typehints.TypeVariable('VT')]),
+        (
+            'bare tuple',
+            tuple,
+            typehints.Tuple[typehints.TypeVariable('T'), ...]),
+        ('bare set', typing.Set, typehints.Set[typehints.TypeVariable('T')]),
+        (
+            'bare frozenset',
+            frozenset,
+            typehints.FrozenSet[typehints.TypeVariable(
+                'T', use_name_in_eq=False)]),
+        (
+            'bare iterator',
+            collections.abc.Iterator,
+            typehints.Iterator[typehints.TypeVariable('T_co')]),
+        (
+            'bare iterable',
+            collections.abc.Iterable,
+            typehints.Iterable[typehints.TypeVariable('T_co')]),
+        (
+            'nested bare',
+            tuple[collections.abc.Iterator],
+            typehints.Tuple[typehints.Iterator[typehints.TypeVariable('T_co')]]
+        ),
+        (
+            'bare generator',
+            collections.abc.Generator,
+            typehints.Generator[typehints.TypeVariable('T_co')]),
+    ]
+    for test_case in test_cases:
+      description = test_case[0]
+      typing_type = test_case[1]
+      expected_beam_type = test_case[2]
+      converted_beam_type = convert_to_beam_type(typing_type)
+      self.assertEqual(expected_beam_type, converted_beam_type, description)
+
   def test_convert_bare_types_fail(self):
     # These conversions should fail.
     test_cases = [
diff --git a/sdks/python/apache_beam/typehints/typehints_test.py b/sdks/python/apache_beam/typehints/typehints_test.py
@@ -1250,6 +1250,26 @@ def test_mix_fails(self):
     self.assertNotCompatible(re.Pattern[str], re.Match[str])
 
 
+# TODO(https://github.com/apache/beam/issues/34644): implement a
+# CallableTypeHint class and implement test suites for inner type checking
+class CallableTestCase(TypeHintTestCase):
+  def test_callable(self):
+    self.assertCompatible(collections.abc.Callable, collections.abc.Callable)
+    self.assertCompatible(
+        collections.abc.Callable[[str], int],
+        collections.abc.Callable[[str], int])
+
+  def test_normalize_to_any_bare(self):
+    input = collections.abc.Callable
+    output = typehints.normalize(input)
+    self.assertEqual(Any, output)
+
+  def test_normalize_to_any_parameterized(self):
+    input = collections.abc.Callable[[str], int]
+    output = typehints.normalize(input)
+    self.assertEqual(Any, output)
+
+
 class TakesDecoratorTestCase(TypeHintTestCase):
   def test_must_be_primitive_type_or_constraint(self):
     with self.assertRaises(TypeError) as e:
