dycw-utilities 0.134.4__tar.gz → 0.135.0__tar.gz

{dycw_utilities-0.134.4 → dycw_utilities-0.135.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.134.4
+Version: 0.135.0
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12

{dycw_utilities-0.134.4 → dycw_utilities-0.135.0}/pyproject.toml

@@ -56,7 +56,7 @@ operator = ["polars-lts-cpu"]
 optuna = ["optuna >= 4.4.0, < 4.5"]
 orjson = ["orjson >= 3.10.18, < 3.11"]
 orjson-test = ["polars-lts-cpu"]
-polars = ["polars-lts-cpu >= 1.30.0, < 1.31"]
+polars = ["polars >= 1.30.0, < 1.31", "polars-lts-cpu >= 1.30.0, < 1.31"]
 polars-ols = ["polars-ols >= 0.3.5, < 0.4"]
 polars-ols-test = ["scikit-learn"]
 polars-test = ["numpy", "statsmodels"]
@@ -105,7 +105,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.134.4"
+version = "0.135.0"
 
 [project.optional-dependencies]
 logging = [
@@ -132,7 +132,7 @@ test = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.134.4"
+current_version = "0.135.0"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.134.4 → dycw_utilities-0.135.0}/src/tests/test_more_itertools.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Iterator
+from collections.abc import Collection, Iterator
 from re import escape
 from typing import TYPE_CHECKING, Any, ClassVar, TypeGuard
 
@@ -23,41 +23,74 @@ if TYPE_CHECKING:
 
 class TestBucketMapping:
     iterable: ClassVar[list[str]] = ["a1", "b1", "c1", "a2", "b2", "c2", "b3"]
+    iterable_unique: ClassVar[list[str]] = ["a1", "b2", "c3"]
 
     def test_main(self) -> None:
         mapping = bucket_mapping(self.iterable, lambda x: x[0])
         assert set(mapping) == {"a", "b", "c"}
         for value in mapping.values():
             assert isinstance(value, Iterator)
-            assert not isinstance(value, list)
+            assert not isinstance(value, Collection)
         assert list(mapping["a"]) == ["a1", "a2"]
         assert list(mapping["b"]) == ["b1", "b2", "b3"]
         assert list(mapping["c"]) == ["c1", "c2"]
 
-    def test_transform(self) -> None:
+    def test_list(self) -> None:
+        mapping = bucket_mapping(self.iterable, lambda x: x[0], post="list")
+        assert set(mapping) == {"a", "b", "c"}
+        for value in mapping.values():
+            assert isinstance(value, list)
+        assert mapping["a"] == ["a1", "a2"]
+        assert mapping["b"] == ["b1", "b2", "b3"]
+        assert mapping["c"] == ["c1", "c2"]
+
+    def test_tuple(self) -> None:
+        mapping = bucket_mapping(self.iterable, lambda x: x[0], post="tuple")
+        assert set(mapping) == {"a", "b", "c"}
+        for value in mapping.values():
+            assert isinstance(value, tuple)
+        assert mapping["a"] == ("a1", "a2")
+        assert mapping["b"] == ("b1", "b2", "b3")
+        assert mapping["c"] == ("c1", "c2")
+
+    def test_set(self) -> None:
+        mapping = bucket_mapping(self.iterable, lambda x: x[0], post="set")
+        assert set(mapping) == {"a", "b", "c"}
+        for value in mapping.values():
+            assert isinstance(value, set)
+        assert mapping["a"] == {"a1", "a2"}
+        assert mapping["b"] == {"b1", "b2", "b3"}
+        assert mapping["c"] == {"c1", "c2"}
+
+    def test_frozenset(self) -> None:
+        mapping = bucket_mapping(self.iterable, lambda x: x[0], post="frozenset")
+        assert set(mapping) == {"a", "b", "c"}
+        for value in mapping.values():
+            assert isinstance(value, frozenset)
+        assert mapping["a"] == frozenset({"a1", "a2"})
+        assert mapping["b"] == frozenset({"b1", "b2", "b3"})
+        assert mapping["c"] == frozenset({"c1", "c2"})
+
+    def test_unique(self) -> None:
+        mapping = bucket_mapping(self.iterable_unique, lambda x: x[0], post="unique")
+        expected = {"a": "a1", "b": "b2", "c": "c3"}
+        assert mapping == expected
+
+    def test_pre(self) -> None:
         mapping = bucket_mapping(
-            self.iterable, lambda x: x[0], transform=lambda x: int(x[-1])
+            self.iterable, lambda x: x[0], pre=lambda x: int(x[-1])
         )
         assert set(mapping) == {"a", "b", "c"}
         for value in mapping.values():
             assert isinstance(value, Iterator)
-            assert not isinstance(value, list)
+            assert not isinstance(value, Collection)
         assert list(mapping["a"]) == [1, 2]
         assert list(mapping["b"]) == [1, 2, 3]
         assert list(mapping["c"]) == [1, 2]
 
-    def test_list(self) -> None:
-        mapping = bucket_mapping(self.iterable, lambda x: x[0], list=True)
-        assert set(mapping) == {"a", "b", "c"}
-        for value in mapping.values():
-            assert isinstance(value, list)
-        assert mapping["a"] == ["a1", "a2"]
-        assert mapping["b"] == ["b1", "b2", "b3"]
-        assert mapping["c"] == ["c1", "c2"]
-
-    def test_transform_and_list(self) -> None:
+    def test_pre_and_list(self) -> None:
         mapping = bucket_mapping(
-            self.iterable, lambda x: x[0], transform=lambda x: int(x[-1]), list=True
+            self.iterable, lambda x: x[0], pre=lambda x: int(x[-1]), post="list"
         )
         assert set(mapping) == {"a", "b", "c"}
         for value in mapping.values():
@@ -66,19 +99,45 @@ class TestBucketMapping:
         assert mapping["b"] == [1, 2, 3]
         assert mapping["c"] == [1, 2]
 
-    iterable_unique: ClassVar[list[str]] = ["a1", "b2", "c3"]
+    def test_pre_and_tuple(self) -> None:
+        mapping = bucket_mapping(
+            self.iterable, lambda x: x[0], pre=lambda x: int(x[-1]), post="tuple"
+        )
+        assert set(mapping) == {"a", "b", "c"}
+        for value in mapping.values():
+            assert isinstance(value, tuple)
+        assert mapping["a"] == (1, 2)
+        assert mapping["b"] == (1, 2, 3)
+        assert mapping["c"] == (1, 2)
 
-    def test_unique(self) -> None:
-        mapping = bucket_mapping(self.iterable_unique, lambda x: x[0], unique=True)
-        expected = {"a": "a1", "b": "b2", "c": "c3"}
-        assert mapping == expected
+    def test_pre_and_set(self) -> None:
+        mapping = bucket_mapping(
+            self.iterable, lambda x: x[0], pre=lambda x: int(x[-1]), post="set"
+        )
+        assert set(mapping) == {"a", "b", "c"}
+        for value in mapping.values():
+            assert isinstance(value, set)
+        assert mapping["a"] == {1, 2}
+        assert mapping["b"] == {1, 2, 3}
+        assert mapping["c"] == {1, 2}
+
+    def test_pre_and_frozenset(self) -> None:
+        mapping = bucket_mapping(
+            self.iterable, lambda x: x[0], pre=lambda x: int(x[-1]), post="frozenset"
+        )
+        assert set(mapping) == {"a", "b", "c"}
+        for value in mapping.values():
+            assert isinstance(value, frozenset)
+        assert mapping["a"] == frozenset({1, 2})
+        assert mapping["b"] == frozenset({1, 2, 3})
+        assert mapping["c"] == frozenset({1, 2})
 
-    def test_transform_and_unique(self) -> None:
+    def test_pre_and_unique(self) -> None:
         mapping = bucket_mapping(
             self.iterable_unique,
             lambda x: x[0],
-            transform=lambda x: int(x[-1]),
-            unique=True,
+            pre=lambda x: int(x[-1]),
+            post="unique",
         )
         expected = {"a": 1, "b": 2, "c": 3}
         assert mapping == expected
@@ -90,9 +149,9 @@ class TestBucketMapping:
                 "Buckets must contain exactly one item each; got 'a' (#1: 'a1', #2: 'a2'), 'b' (#1: 'b1', #2: 'b2'), 'c' (#1: 'c1', #2: 'c2')"
             ),
         ):
-            _ = bucket_mapping(self.iterable, lambda x: x[0], unique=True)
+            _ = bucket_mapping(self.iterable, lambda x: x[0], post="unique")
 
-    def test_error_transform_and_unique(self) -> None:
+    def test_error_pre_and_unique(self) -> None:
         with raises(
             BucketMappingError,
             match=escape(
@@ -100,10 +159,7 @@ class TestBucketMapping:
             ),
         ):
             _ = bucket_mapping(
-                self.iterable,
-                lambda x: x[0],
-                transform=lambda x: int(x[-1]),
-                unique=True,
+                self.iterable, lambda x: x[0], pre=lambda x: int(x[-1]), post="unique"
             )
 
 

{dycw_utilities-0.134.4 → dycw_utilities-0.135.0}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.134.4"
+__version__ = "0.135.0"

{dycw_utilities-0.134.4 → dycw_utilities-0.135.0}/src/utilities/more_itertools.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-import builtins
-from collections.abc import Hashable
+from collections.abc import Callable, Hashable
 from dataclasses import dataclass
 from itertools import islice
 from textwrap import indent
@@ -25,7 +24,7 @@ from utilities.reprlib import get_repr
 from utilities.sentinel import Sentinel, sentinel
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
+    from collections.abc import Iterable, Iterator, Mapping, Sequence
 
 
 @overload
@@ -34,55 +33,107 @@ def bucket_mapping[T, U, UH: Hashable](
     func: Callable[[T], UH],
     /,
     *,
-    transform: Callable[[T], U],
-    list: bool = False,
-    unique: Literal[True],
-) -> Mapping[UH, U]: ...
+    pre: Callable[[T], U],
+    post: Literal["list"],
+) -> Mapping[UH, list[U]]: ...
 @overload
 def bucket_mapping[T, U, UH: Hashable](
     iterable: Iterable[T],
     func: Callable[[T], UH],
     /,
     *,
-    transform: Callable[[T], U] | None = None,
-    list: bool = False,
-    unique: Literal[True],
-) -> Mapping[UH, T]: ...
+    pre: Callable[[T], U],
+    post: Literal["tuple"],
+) -> Mapping[UH, tuple[U, ...]]: ...
 @overload
 def bucket_mapping[T, U, UH: Hashable](
     iterable: Iterable[T],
     func: Callable[[T], UH],
     /,
     *,
-    transform: Callable[[T], U],
-    list: Literal[True],
-) -> Mapping[UH, Sequence[U]]: ...
+    pre: Callable[[T], U],
+    post: Literal["set"],
+) -> Mapping[UH, set[U]]: ...
 @overload
 def bucket_mapping[T, U, UH: Hashable](
     iterable: Iterable[T],
     func: Callable[[T], UH],
     /,
     *,
-    transform: Callable[[T], U],
-    list: bool = False,
-) -> Mapping[UH, Iterator[U]]: ...
+    pre: Callable[[T], U],
+    post: Literal["frozenset"],
+) -> Mapping[UH, frozenset[U]]: ...
 @overload
 def bucket_mapping[T, U, UH: Hashable](
     iterable: Iterable[T],
     func: Callable[[T], UH],
     /,
     *,
-    transform: Callable[[T], U] | None = None,
-    list: Literal[True],
-) -> Mapping[UH, Sequence[T]]: ...
+    pre: Callable[[T], U] | None = None,
+    post: Literal["unique"],
+) -> Mapping[UH, U]: ...
 @overload
 def bucket_mapping[T, U, UH: Hashable](
     iterable: Iterable[T],
     func: Callable[[T], UH],
     /,
     *,
-    transform: Callable[[T], U] | None = None,
-    list: bool = False,
+    pre: Callable[[T], U] | None = None,
+    post: None = None,
+) -> Mapping[UH, Iterator[U]]: ...
+@overload
+def bucket_mapping[T, UH: Hashable](
+    iterable: Iterable[T],
+    func: Callable[[T], UH],
+    /,
+    *,
+    pre: None = None,
+    post: Literal["list"],
+) -> Mapping[UH, list[T]]: ...
+@overload
+def bucket_mapping[T, UH: Hashable](
+    iterable: Iterable[T],
+    func: Callable[[T], UH],
+    /,
+    *,
+    pre: None = None,
+    post: Literal["tuple"],
+) -> Mapping[UH, tuple[T, ...]]: ...
+@overload
+def bucket_mapping[T, UH: Hashable](
+    iterable: Iterable[T],
+    func: Callable[[T], UH],
+    /,
+    *,
+    pre: None = None,
+    post: Literal["set"],
+) -> Mapping[UH, set[T]]: ...
+@overload
+def bucket_mapping[T, UH: Hashable](
+    iterable: Iterable[T],
+    func: Callable[[T], UH],
+    /,
+    *,
+    pre: None = None,
+    post: Literal["frozenset"],
+) -> Mapping[UH, frozenset[T]]: ...
+@overload
+def bucket_mapping[T, UH: Hashable](
+    iterable: Iterable[T],
+    func: Callable[[T], UH],
+    /,
+    *,
+    pre: None = None,
+    post: Literal["unique"],
+) -> Mapping[UH, T]: ...
+@overload
+def bucket_mapping[T, UH: Hashable](
+    iterable: Iterable[T],
+    func: Callable[[T], UH],
+    /,
+    *,
+    pre: None = None,
+    post: None = None,
 ) -> Mapping[UH, Iterator[T]]: ...
 @overload
 def bucket_mapping[T, U, UH: Hashable](
@@ -90,15 +141,20 @@ def bucket_mapping[T, U, UH: Hashable](
     func: Callable[[T], UH],
     /,
     *,
-    transform: Callable[[T], U] | None = None,
-    list: bool = False,
-    unique: bool = False,
+    pre: Callable[[T], U] | None = None,
+    post: Literal["list", "tuple", "set", "frozenset", "unique"] | None = None,
 ) -> (
     Mapping[UH, Iterator[T]]
-    | Mapping[UH, Iterator[U]]
-    | Mapping[UH, Sequence[T]]
-    | Mapping[UH, Sequence[U]]
+    | Mapping[UH, list[T]]
+    | Mapping[UH, tuple[T, ...]]
+    | Mapping[UH, set[T]]
+    | Mapping[UH, frozenset[T]]
     | Mapping[UH, T]
+    | Mapping[UH, Iterator[U]]
+    | Mapping[UH, list[U]]
+    | Mapping[UH, tuple[U, ...]]
+    | Mapping[UH, set[U]]
+    | Mapping[UH, frozenset[U]]
     | Mapping[UH, U]
 ): ...
 def bucket_mapping[T, U, UH: Hashable](
@@ -106,42 +162,66 @@ def bucket_mapping[T, U, UH: Hashable](
     func: Callable[[T], UH],
     /,
     *,
-    transform: Callable[[T], U] | None = None,
-    list: bool = False,  # noqa: A002
-    unique: bool = False,
+    pre: Callable[[T], U] | None = None,
+    post: Literal["list", "tuple", "set", "frozenset", "unique"] | None = None,
 ) -> (
     Mapping[UH, Iterator[T]]
-    | Mapping[UH, Iterator[U]]
-    | Mapping[UH, Sequence[T]]
-    | Mapping[UH, Sequence[U]]
+    | Mapping[UH, list[T]]
+    | Mapping[UH, tuple[T, ...]]
+    | Mapping[UH, set[T]]
+    | Mapping[UH, frozenset[T]]
     | Mapping[UH, T]
+    | Mapping[UH, Iterator[U]]
+    | Mapping[UH, list[U]]
+    | Mapping[UH, tuple[U, ...]]
+    | Mapping[UH, set[U]]
+    | Mapping[UH, frozenset[U]]
     | Mapping[UH, U]
 ):
     """Bucket the values of iterable into a mapping."""
-    b = bucket(iterable, func)
-    mapping = {key: b[key] for key in b}
-    match transform, list:
-        case None, False:
-            ...
-        case None, True:
-            mapping = {k: builtins.list(v) for k, v in mapping.items()}
-        case _, False:
-            mapping = {k: map(transform, v) for k, v in mapping.items()}
-        case _, True:
-            mapping = {k: builtins.list(map(transform, v)) for k, v in mapping.items()}
+    bckt = bucket(iterable, func)
+    mapping = {key: bckt[key] for key in bckt}
+    match pre, post:
+        case None, None:
+            return mapping
+        case None, "list":
+            return {k: list(v) for k, v in mapping.items()}
+        case None, "tuple":
+            return {k: tuple(v) for k, v in mapping.items()}
+        case None, "set":
+            return {k: set(v) for k, v in mapping.items()}
+        case None, "frozenset":
+            return {k: frozenset(v) for k, v in mapping.items()}
+        case None, "unique":
+            return _bucket_mapping_unique(mapping)
+        case Callable(), None:
+            return {k: map(pre, v) for k, v in mapping.items()}
+        case Callable(), "list":
+            return {k: list(map(pre, v)) for k, v in mapping.items()}
+        case Callable(), "tuple":
+            return {k: tuple(map(pre, v)) for k, v in mapping.items()}
+        case Callable(), "set":
+            return {k: set(map(pre, v)) for k, v in mapping.items()}
+        case Callable(), "frozenset":
+            return {k: frozenset(map(pre, v)) for k, v in mapping.items()}
+        case Callable(), "unique":
+            return _bucket_mapping_unique({k: map(pre, v) for k, v in mapping.items()})
         case _ as never:
             assert_never(never)
-    if not unique:
-        return mapping
-    results = {}
-    error_no_transform: dict[UH, tuple[T, T]] = {}
+
+
+def _bucket_mapping_unique[K: Hashable, V](
+    mapping: Mapping[K, Iterable[V]], /
+) -> Mapping[K, V]:
+    results: dict[K, V] = {}
+    errors: dict[K, tuple[V, V]] = {}
     for key, value in mapping.items():
         try:
             results[key] = one(value)
         except OneNonUniqueError as error:
-            error_no_transform[key] = (error.first, error.second)
-    if len(error_no_transform) >= 1:
-        raise BucketMappingError(errors=error_no_transform)
+            errors[key] = (error.first, error.second)
+    if len(errors) >= 1:
+        raise BucketMappingError(errors=errors)
     return results