def test_iterables(self):
self.assertEqual(int, typehints.get_yielded_type(typehints.Iterable[int]))
self.assertEqual(int, typehints.get_yielded_type(typehints.Iterator[int]))
self.assertEqual(int, typehints.get_yielded_type(typehints.Generator[int]))
self.assertEqual(int, typehints.get_yielded_type(typehints.List[int]))
self.assertEqual(
typehints.Union[int, str],
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]))
def strip_iterable(self):
"""Removes outer Iterable (or equivalent) from output type.
Only affects instances with simple output types, otherwise is a no-op.
Example: Generator[Tuple(int, int)] becomes Tuple(int, int)
Raises:
ValueError if output type is simple and not iterable.
"""
if not self.has_simple_output_type():
return
yielded_type = typehints.get_yielded_type(self.output_types[0][0])
self.output_types = ((yielded_type, ), {})
def strip_iterable(self):
"""Removes outer Iterable (or equivalent) from output type.
Only affects instances with simple output types, otherwise is a no-op.
Does not modify self.
Example: Generator[Tuple(int, int)] becomes Tuple(int, int)
Returns:
A possible copy of this instance with a possibly different output type.
Raises:
ValueError if output type is simple and not iterable.
"""
if not self.has_simple_output_type():
return self
yielded_type = typehints.get_yielded_type(self.output_types[0][0])
res = self.copy()
res.output_types = ((yielded_type,), {})
return res
def strip_iterable(self):
# type: () -> IOTypeHints
"""Removes outer Iterable (or equivalent) from output type.
Only affects instances with simple output types, otherwise is a no-op.
Does not modify self.
Designed to be used with type hints from callables of ParDo, FlatMap, DoFn.
Output type may be Optional[T], in which case the result of stripping T is
used as the output type.
Output type may be None/NoneType, in which case nothing is done.
Example: Generator[Tuple(int, int)] becomes Tuple(int, int)
Returns:
A copy of this instance with a possibly different output type.
Raises:
ValueError if output type is simple and not iterable.
"""
if self.output_types is None or not self.has_simple_output_type():
return self
output_type = self.output_types[0][0]
if output_type is None or isinstance(output_type, type(None)):
return self
# If output_type == Optional[T]: output_type = T.
if isinstance(output_type, typehints.UnionConstraint):
types = list(output_type.union_types)
if len(types) == 2:
try:
types.remove(type(None))
output_type = types[0]
except ValueError:
pass
yielded_type = typehints.get_yielded_type(output_type)
return self._replace(output_types=((yielded_type, ), {}),
origin=self._make_origin([self],
tb=False,
msg=['strip_iterable()'
]))
def test_not_iterable(self):
with self.assertRaisesRegex(ValueError, r'not iterable'):
typehints.get_yielded_type(int)
