From bbcaed3314507e74acc380f31b24d924451e7acf Mon Sep 17 00:00:00 2001 From: Derek Wan Date: Wed, 9 Oct 2024 20:28:25 +0900 Subject: [PATCH] extend --- src/tests/test_iterables.py | 12 ++++++++++++ src/utilities/iterables.py | 5 +++++ 2 files changed, 17 insertions(+) diff --git a/src/tests/test_iterables.py b/src/tests/test_iterables.py index 0c48300a7..40b44d0f2 100644 --- a/src/tests/test_iterables.py +++ b/src/tests/test_iterables.py @@ -592,6 +592,18 @@ def test_new_many_iterables(self) -> None: assert isinstance(collection, Collection) assert len(collection) == 3 + def test_new_post_init(self) -> None: + class SubCollection(Collection[_Item]): + @classmethod + @override + def check_items(cls, items: Iterable[_Item]) -> None: + if any(item.n >= 1 for item in items): + msg = "n >= 1 is not permitted" + raise ValueError(msg) + + with raises(ValueError, match="n >= 1 is not permitted"): + _ = SubCollection(map(_Item, range(3))) + def test_or_singleton(self) -> None: collection = Collection(map(_Item, range(3))) result = collection | _Item(3) diff --git a/src/utilities/iterables.py b/src/utilities/iterables.py index 5b729e9fc..fc002a01a 100644 --- a/src/utilities/iterables.py +++ b/src/utilities/iterables.py @@ -499,6 +499,7 @@ class Collection(frozenset[_TSupportsHashAndSort]): def __new__(cls, *item_or_items: MaybeIterable[_TSupportsHashAndSort]) -> Self: items = list(chain(*map(always_iterable, item_or_items))) + cls.check_items(items) return super().__new__(cls, items) def __init__(self, *item_or_items: MaybeIterable[_TSupportsHashAndSort]) -> None: @@ -574,6 +575,10 @@ def __sub__( other = cast(Iterable[_TSupportsHashAndSort], other) return self.__sub__(type(self)(other)) + @classmethod + def check_items(cls, items: Iterable[_TSupportsHashAndSort], /) -> None: + _ = items + def filter(self, func: Callable[[_TSupportsHashAndSort], bool], /) -> Self: return type(self)(filter(func, self))