Extend partition to accept multiple predicates

Author: kurtbrose

boltons/iterutils.py

@@ -754,10 +754,15 @@ def value_transform(x): return f(x[1])
     return ret
 
 
-def partition(src, key=bool):
+def partition(src, key=bool, *keys):
     """No relation to :meth:`str.partition`, ``partition`` is like
-    :func:`bucketize`, but for added convenience returns a tuple of
-    ``(truthy_values, falsy_values)``.
+    :func:`bucketize`, but for added convenience returns a collection for
+    each predicate passed.
+
+    ``partition`` now accepts multiple *key* functions and will return
+    ``N + 1`` lists for ``N`` predicates. Each value from *src* is placed
+    into the first list whose predicate evaluates to ``True`` with values
+    that match none of the predicates placed in the last list.
 
     >>> nonempty, empty = partition(['', '', 'hi', '', 'bye'])
     >>> nonempty
@@ -772,9 +777,37 @@ def partition(src, key=bool):
     >>> decimal_digits, hexletters = partition(string.hexdigits, is_digit)
     >>> ''.join(decimal_digits), ''.join(hexletters)
     ('0123456789', 'abcdefABCDEF')
+
+    Multiple predicates may be supplied to divide into more buckets:
+
+    >>> positive, negative, zero = partition(range(-1, 2),
+    ...                                     lambda i: i > 0,
+    ...                                     lambda i: i < 0)
+    >>> positive, negative, zero
+    ([1], [-1], [0])
     """
-    bucketized = bucketize(src, key)
-    return bucketized.get(True, []), bucketized.get(False, [])
+    if not is_iterable(src):
+        raise TypeError('expected an iterable')
+
+    def _make_key_func(k):
+        if isinstance(k, str):
+            return lambda x, k=k: getattr(x, k, False)
+        if callable(k):
+            return k
+        raise TypeError('expected key to be callable or a string')
+
+    key_funcs = [_make_key_func(key)] + [_make_key_func(k) for k in keys]
+    parts = [[] for _ in range(len(key_funcs) + 1)]
+
+    for val in src:
+        for idx, func in enumerate(key_funcs):
+            if func(val):
+                parts[idx].append(val)
+                break
+        else:
+            parts[-1].append(val)
+
+    return tuple(parts)
 
 
 def unique(src, key=None):

tests/test_iterutils.py

@@ -616,3 +616,22 @@ def test_windowed_filled():
 
     assert list(windowed_iter(range(4), 3)) == [(0, 1, 2), (1, 2, 3)]
     assert list(windowed_iter(range(4), 3, fill=None)) == [(0, 1, 2), (1, 2, 3), (2, 3, None), (3, None, None)]
+# Tests for partition
+
+def test_partition_default():
+    from boltons.iterutils import partition
+
+    nonempty, empty = partition(['', '', 'hi', '', 'bye'])
+    assert nonempty == ['hi', 'bye']
+    assert empty == ["", "", ""]
+
+
+def test_partition_multiple():
+    from boltons.iterutils import partition
+
+    items = [2, -1, 0, 3, -2]
+    positive, negative, zero = partition(items, lambda i: i > 0, lambda i: i < 0)
+    assert positive == [2, 3]
+    assert negative == [-1, -2]
+    assert zero == [0]
+