feat: block partition (#79)

* feat: block partition

* chore: update pyright

* fix: export partition

Author: dbrattli

.pre-commit-config.yaml

@@ -28,7 +28,7 @@ repos:
         language: node
         pass_filenames: false
         types: [python]
-        additional_dependencies: ["[email protected].241"]
+        additional_dependencies: ["[email protected].243"]
     repo: local
   - hooks:
     - id: jb-to-sphinx

expression/collections/block.py

@@ -47,6 +47,7 @@
     Some,
     SupportsLessThan,
     SupportsSum,
+    curry_flipped,
     pipe,
 )
 from expression.core.typing import GenericValidator, ModelField, SupportsValidation
@@ -364,6 +365,34 @@ def of_seq(xs: Iterable[_TSource]) -> Block[_TSource]:
     def of_option(option: Option[_TSource]) -> Block[_TSource]:
         return of_option(option)
 
+    def partition(
+        self, predicate: Callable[[_TSource], bool]
+    ) -> Tuple[Block[_TSource], Block[_TSource]]:
+        """Partition block.
+
+        Splits the collection into two collections, containing the
+        elements for which the given predicate returns True and False
+        respectively. Element order is preserved in both of the created
+        lists.
+
+        Args:
+            predicate: The function to test the input elements.
+
+        Returns:
+            A list containing the elements for which the predicate
+            evaluated to true and a list containing the elements for
+            which the predicate evaluated to false.
+        """
+        list1: List[_TSource] = []
+        list2: List[_TSource] = []
+
+        for item in self.value:
+            if predicate(item):
+                list1.append(item)
+            else:
+                list2.append(item)
+        return (Block(list1), Block(list2))
+
     @overload
     @staticmethod
     def range(stop: int) -> Block[int]:
@@ -580,9 +609,10 @@ def _choose(source: Block[_TSource]) -> Block[_TResult]:
     return _choose
 
 
+@curry_flipped(1)
 def collect(
-    mapping: Callable[[_TSource], Block[_TResult]]
-) -> Callable[[Block[_TSource]], Block[_TResult]]:
+    source: Block[_TSource], mapping: Callable[[_TSource], Block[_TResult]]
+) -> Block[_TResult]:
     """For each element of the list, applies the given function.
     Concatenates all the results and return the combined list.
 
@@ -595,19 +625,7 @@ def collect(
         list and returns the concatenation of the transformed sublists.
     """
 
-    def _collect(source: Block[_TSource]) -> Block[_TResult]:
-        """For each element of the list, applies the given function.
-        Concatenates all the results and return the combined list.
-
-        Args:
-            source: The input list.
-
-        Returns:
-            The concatenation of the transformed sublists.
-        """
-        return source.collect(mapping)
-
-    return _collect
+    return source.collect(mapping)
 
 
 def concat(sources: Iterable[Block[_TSource]]) -> Block[_TSource]:
@@ -865,6 +883,29 @@ def of_option(option: Option[_TSource]) -> Block[_TSource]:
     return empty
 
 
+@curry_flipped(1)
+def partition(
+    source: Block[_TSource], predicate: Callable[[_TSource], bool]
+) -> Tuple[Block[_TSource], Block[_TSource]]:
+    """Partition block.
+
+    Splits the collection into two collections, containing the
+    elements for which the given predicate returns True and False
+    respectively. Element order is preserved in both of the created
+    lists.
+
+    Args:
+        source: The source block to partition (curried flipped)
+        predicate: The function to test the input elements.
+
+    Returns:
+        A list containing the elements for which the predicate
+        evaluated to true and a list containing the elements for
+        which the predicate evaluated to false.
+    """
+    return source.partition(predicate)
+
+
 @overload
 def range(stop: int) -> Block[int]:
     ...
@@ -1104,6 +1145,7 @@ def _zip(source: Block[_TSource]) -> Block[Tuple[_TSource, _TResult]]:
     "choose",
     "collect",
     "concat",
+    "dict",
     "empty",
     "filter",
     "fold",
@@ -1115,6 +1157,7 @@ def _zip(source: Block[_TSource]) -> Block[Tuple[_TSource, _TResult]]:
     "mapi",
     "of_seq",
     "of_option",
+    "partition",
     "singleton",
     "skip",
     "skip_last",
@@ -1123,7 +1166,6 @@ def _zip(source: Block[_TSource]) -> Block[Tuple[_TSource, _TResult]]:
     "tail",
     "take",
     "take_last",
-    "dict",
     "try_head",
     "unfold",
     "zip",