Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for collections.abc.Collection type hints #29272

Merged
merged 5 commits into from
Nov 7, 2023
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
_CONVERTED_COLLECTIONS = [
collections.abc.Set,
collections.abc.MutableSet,
collections.abc.Collection,
]


Expand Down Expand Up @@ -118,6 +119,10 @@ def _match_is_exactly_iterable(user_type):
return getattr(user_type, '__origin__', None) is expected_origin


def _match_is_exactly_collection(user_type):
return getattr(user_type, '__origin__', None) is collections.abc.Collection


def match_is_named_tuple(user_type):
return (
_safe_issubclass(user_type, typing.Tuple) and
Expand Down Expand Up @@ -322,6 +327,10 @@ def convert_to_beam_type(typ):
match=_match_issubclass(typing.Iterator),
arity=1,
beam_type=typehints.Iterator),
_TypeMapEntry(
match=_match_is_exactly_collection,
arity=1,
beam_type=typehints.Collection),
]

# Find the first matching entry.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,15 @@ def test_convert_to_beam_type_with_collections_types(self):
(
'enum mutable set',
collections.abc.MutableSet[_TestEnum],
typehints.Set[_TestEnum])
typehints.Set[_TestEnum]),
(
'collection enum',
collections.abc.Collection[_TestEnum],
typehints.Collection[_TestEnum]),
(
'collection of tuples',
collections.abc.Collection[tuple[str, int]],
typehints.Collection[typehints.Tuple[str, int]]),
]

for test_case in test_cases:
Expand Down
54 changes: 54 additions & 0 deletions sdks/python/apache_beam/typehints/typehints.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,7 @@
'Dict',
'Set',
'FrozenSet',
'Collection',
'Iterable',
'Iterator',
'Generator',
Expand Down Expand Up @@ -1017,6 +1018,58 @@ def __getitem__(self, type_param):
FrozenSetTypeConstraint = FrozenSetHint.FrozenSetTypeConstraint


class CollectionHint(CompositeTypeHint):
""" A Collection type-hint.

Collection[X] defines a type-hint for a collection of homogenous types. 'X'
may be either a built-in Python type or another nested TypeConstraint.

This represents a collections.abc.Collection type, which implements
__contains__, __iter__, and __len__. This acts as a parent type for
sets but has fewer guarantees for mixins.
"""
class CollectionTypeConstraint(SequenceTypeConstraint):
def __init__(self, type_param):
super().__init__(type_param, abc.Collection)

def __repr__(self):
return 'Collection[%s]' % repr(self.inner_type)

@staticmethod
def _is_subclass_constraint(sub):
return isinstance(
sub, (
CollectionTypeConstraint,
FrozenSetTypeConstraint,
SetTypeConstraint))

def _consistent_with_check_(self, sub):
if self._is_subclass_constraint(sub):
return is_consistent_with(sub.inner_type, self.inner_type)
elif isinstance(sub, TupleConstraint):
if not sub.tuple_types:
# The empty tuple is consistent with Iterator[T] for any T.
return True
# Each element in the hetrogenious tuple must be consistent with
# the iterator type.
# E.g. Tuple[A, B] < Iterable[C] if A < C and B < C.
jrmccluskey marked this conversation as resolved.
Show resolved Hide resolved
return all(
is_consistent_with(elem, self.inner_type)
for elem in sub.tuple_types)
elif not isinstance(sub, TypeConstraint):
if getattr(sub, '__origin__', None) is not None:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need this check? I would've assumed issubclass would take care of it

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Parameterized generics and issubclass do not get along. If you pass along a parameterized type (something like issubclass(set[int], collections.abc.Collection)) you actually get an error that the first arg has to be a class. This is why the internal checking for these containers separates out the container type from the type it wraps and checks them separately.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Got it - thanks

return issubclass(sub, abc.Collection)
return False

def __getitem__(self, type_param):
validate_composite_type_param(
type_param, error_msg_prefix='Parameter to a Collection hint')
return self.CollectionTypeConstraint(type_param)


CollectionTypeConstraint = CollectionHint.CollectionTypeConstraint


class IterableHint(CompositeTypeHint):
"""An Iterable type-hint.

Expand Down Expand Up @@ -1187,6 +1240,7 @@ def __getitem__(self, type_params):
Dict = DictHint()
Set = SetHint()
FrozenSet = FrozenSetHint()
Collection = CollectionHint()
Iterable = IterableHint()
Iterator = IteratorHint()
Generator = GeneratorHint()
Expand Down
27 changes: 27 additions & 0 deletions sdks/python/apache_beam/typehints/typehints_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -865,6 +865,32 @@ class FrozenSetHintTestCase(BaseSetHintTest.CommonTests):
string_type = 'FrozenSet'


class CollectionHintTestCase(TypeHintTestCase):
def test_type_constraint_compatibility(self):
self.assertCompatible(typehints.Collection[int], typehints.Set[int])
self.assertCompatible(typehints.Iterable[int], typehints.Collection[int])
self.assertCompatible(typehints.Collection[int], typehints.FrozenSet[int])
self.assertCompatible(
typehints.Collection[typehints.Any], typehints.Collection[int])
self.assertCompatible(typehints.Collection[int], typehints.Tuple[int])

def test_one_way_compatibility(self):
self.assertNotCompatible(typehints.Set[int], typehints.Collection[int])
self.assertNotCompatible(
typehints.FrozenSet[int], typehints.Collection[int])
self.assertNotCompatible(typehints.Tuple[int], typehints.Collection[int])
self.assertNotCompatible(typehints.Collection[int], typehints.Iterable[int])

def test_getitem_invalid_composite_type_param(self):
with self.assertRaises(TypeError) as e:
typehints.Collection[5]
self.assertEqual(
'Parameter to a Collection hint must be a '
'non-sequence, a type, or a TypeConstraint. 5 is '
'an instance of int.',
e.exception.args[0])


class IterableHintTestCase(TypeHintTestCase):
def test_getitem_invalid_composite_type_param(self):
with self.assertRaises(TypeError) as e:
Expand Down Expand Up @@ -893,6 +919,7 @@ def test_compatibility(self):
self.assertCompatible(
typehints.Iterable[typehints.Any],
typehints.List[typehints.Tuple[int, bool]])
self.assertCompatible(typehints.Iterable[int], typehints.Collection[int])

def test_tuple_compatibility(self):
self.assertCompatible(typehints.Iterable[int], typehints.Tuple[int, ...])
Expand Down
Loading