diff --git a/sdks/python/apache_beam/typehints/typehints.py b/sdks/python/apache_beam/typehints/typehints.py index ffef40de6673..041a719ff831 100644 --- a/sdks/python/apache_beam/typehints/typehints.py +++ b/sdks/python/apache_beam/typehints/typehints.py @@ -1569,7 +1569,7 @@ def regex_consistency(sub, base) -> bool: def get_yielded_type(type_hint): - """Obtains the type of elements yielded by an iterable.s + """Obtains the type of elements yielded by an iterable. Note that "iterable" here means: can be iterated over in a for loop, excluding strings and dicts. @@ -1583,10 +1583,16 @@ def get_yielded_type(type_hint): Raises: ValueError if not iterable. """ + type_hint = normalize(type_hint, none_as_type=True) if isinstance(type_hint, typing.TypeVar): return typing.Any if isinstance(type_hint, AnyTypeConstraint): return type_hint + 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] if is_consistent_with(type_hint, Iterator[Any]): return type_hint.yielded_type if is_consistent_with(type_hint, Tuple[Any, ...]): @@ -1595,11 +1601,6 @@ 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) diff --git a/sdks/python/apache_beam/typehints/typehints_test.py b/sdks/python/apache_beam/typehints/typehints_test.py index b097a01fed42..d33c6789d840 100644 --- a/sdks/python/apache_beam/typehints/typehints_test.py +++ b/sdks/python/apache_beam/typehints/typehints_test.py @@ -1703,6 +1703,22 @@ def test_iterables(self): typehints.Union[int, str], typehints.get_yielded_type( typehints.Union[typehints.List[int], typehints.List[str]])) + self.assertEqual( + typehints.Union[int, str], + typehints.get_yielded_type( + typehints.Union[typehints.Iterator[int], typehints.Iterator[str]])) + self.assertEqual( + typehints.Union[int, str], + typehints.get_yielded_type( + typehints.Union[typehints.Tuple[int, int], + typehints.Tuple[str, str]])) + self.assertEqual( + typehints.Union[int, str], + typehints.get_yielded_type(typing.Iterable[int] | typing.Iterable[str])) + self.assertEqual( + typehints.Union[int, str], + typehints.get_yielded_type( + typing.Union[typing.Iterator[int], typing.Iterator[str]])) def test_not_iterable(self): with self.assertRaisesRegex(ValueError, r'not iterable'): @@ -1711,6 +1727,8 @@ def test_not_iterable(self): def test_union_not_iterable(self): with self.assertRaisesRegex(ValueError, r'not iterable'): typehints.get_yielded_type(typehints.Union[int, typehints.List[int]]) + with self.assertRaisesRegex(ValueError, r'not iterable'): + typehints.get_yielded_type(int | typing.List[str]) class TestCoerceToKvType(TypeHintTestCase):