Catch union-of-iterables case in get_yielded_type() (#34186)

* Catch union-of-iterables case in get_yielded_type()

* add test case for mixed union

Author: jrmccluskey

sdks/python/apache_beam/typehints/typehints.py

@@ -1401,6 +1401,11 @@ def get_yielded_type(type_hint):
     else:  # TupleSequenceConstraint
       return type_hint.inner_type
   if is_consistent_with(type_hint, Iterable[Any]):
+    if isinstance(type_hint, UnionConstraint):
+      yielded_types = set()
+      for typ in type_hint.inner_types():
+        yielded_types.add(get_yielded_type(typ))
+      return Union[yielded_types]
     return type_hint.inner_type
   raise ValueError('%s is not iterable' % type_hint)
 

sdks/python/apache_beam/typehints/typehints_test.py

@@ -1450,11 +1450,19 @@ def test_iterables(self):
         typehints.get_yielded_type(typehints.Tuple[int, str]))
     self.assertEqual(int, typehints.get_yielded_type(typehints.Set[int]))
     self.assertEqual(int, typehints.get_yielded_type(typehints.FrozenSet[int]))
+    self.assertEqual(
+        typehints.Union[int, str],
+        typehints.get_yielded_type(
+            typehints.Union[typehints.List[int], typehints.List[str]]))
 
   def test_not_iterable(self):
     with self.assertRaisesRegex(ValueError, r'not iterable'):
       typehints.get_yielded_type(int)
 
+  def test_union_not_iterable(self):
+    with self.assertRaisesRegex(ValueError, r'not iterable'):
+      typehints.get_yielded_type(typehints.Union[int, typehints.List[int]])
+
 
 class TestCoerceToKvType(TypeHintTestCase):
   def test_coercion_success(self):