Repository: more-itertools/more-itertools (python)

Issue:
Add powerset_of_sets()
The existing `powerset()` recipe yields tuples.  If actual sets are needed, the obvious extension `map(set, powerset(iterable))`, does not perform well because the `__hash__()` method is called `n * 2**(n-1)` times.

A better approach is to create sets of size 1 and them merge them as needed with `set.union` which reuses the hash values cached in the input sets.

```
def powerset_of_sets(iterable):
    "Return sets of a powerset calling hash() only once per element."
    # powerset_of_sets([1,2,3]) --> set() {1} {2} {3} {1,2} {1,3} {2,3} {1,2,3}
    sets = list(map(set, zip(iterable)))
    for r in range(len(sets) + 1):
        yield from starmap(set().union, combinations(sets, r))
```

Demonstration:

```
from fractions import Fraction

class F(Fraction):
    def __hash__(self):
        "Print a dot every time hash is called"
        print('.', end='')
        return super().__hash__()

>>> _ = list(map(set, powerset(map(F, range(5)))))
................................................................................
>>> _ = list(map(set, powerset_of_sets(map(F, range(5)))))
.....
```

Pull request description:
Issue 820: add powerset_of_sets
This PR adds the implementation `powerset_of_sets` from [this comment](https://github.com/more-itertools/more-itertools/issues/820#issuecomment-1999924928).
 
Closes #820

Patch:
diff --git a/docs/api.rst b/docs/api.rst
index 51e8107..a6a32c5 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -234,6 +234,7 @@ These tools yield combinatorial arrangements of items from iterables.
 .. autofunction:: combination_with_replacement_index
 .. autofunction:: gray_product
 .. autofunction:: outer_product
+.. autofunction:: powerset_of_sets
 
 ----
 
diff --git a/more_itertools/more.py b/more_itertools/more.py
index ae46be6..075a71f 100755
--- a/more_itertools/more.py
+++ b/more_itertools/more.py
@@ -6,6 +6,7 @@ from functools import cached_property, partial, reduce, wraps
 from heapq import heapify, heapreplace, heappop
 from itertools import (
     chain,
+    combinations,
     compress,
     count,
     cycle,
@@ -109,6 +110,7 @@ __all__ = [
     'partitions',
     'peekable',
     'permutation_index',
+    'powerset_of_sets',
     'product_index',
     'raise_',
     'repeat_each',
@@ -4689,3 +4691,19 @@ def filter_map(func, iterable):
         y = func(x)
         if y is not None:
             yield y
+
+
+def powerset_of_sets(iterable):
+    """Yields all possible subsets of the iterable.
+
+        >>> list(powerset_of_sets([1, 2, 3]))  # doctest: +SKIP
+        [set(), {1}, {2}, {3}, {1, 2}, {1, 3}, {2, 3}, {1, 2, 3}]
+        >>> list(powerset_of_sets([1, 1, 0]))  # doctest: +SKIP
+        [set(), {1}, {0}, {0, 1}]
+
+    :func:`powerset_of_sets` takes care to minimize the number
+    of hash operations performed.
+    """
+    sets = tuple(map(set, dict.fromkeys(map(frozenset, zip(iterable)))))
+    for r in range(len(sets) + 1):
+        yield from starmap(set().union, combinations(sets, r))
diff --git a/more_itertools/more.pyi b/more_itertools/more.pyi
index 93b9997..c866d55 100644
--- a/more_itertools/more.pyi
+++ b/more_itertools/more.pyi
@@ -694,3 +694,4 @@ def filter_map(
     func: Callable[[_T], _V | None],
     iterable: Iterable[_T],
 ) -> Iterator[_V]: ...
+def powerset_of_sets(iterable: Iterable[_T]) -> Iterator[set[_T]]: ...
diff --git a/more_itertools/recipes.py b/more_itertools/recipes.py
index abbdd5f..e70abc4 100644
--- a/more_itertools/recipes.py
+++ b/more_itertools/recipes.py
@@ -199,7 +199,7 @@ def nth(iterable, n, default=None):
     return next(islice(iterable, n, None), default)
 
 
-def all_equal(iterable):
+def all_equal(iterable, key=None):
     """
     Returns ``True`` if all the elements are equal to each other.
 
@@ -208,9 +208,16 @@ def all_equal(iterable):
         >>> all_equal('aaab')
         False
 
+    A function that accepts a single argument and returns a transformed version
+    of each input item can be specified with *key*:
+
+        >>> all_equal('AaaA', key=str.casefold)
+        True
+        >>> all_equal([1, 2, 3], key=lambda x: x < 10)
+        True
+
     """
-    g = groupby(iterable)
-    return next(g, True) and not next(g, False)
+    return len(list(islice(groupby(iterable, key), 2))) <= 1
 
 
 def quantify(iterable, pred=bool):
@@ -411,16 +418,11 @@ def roundrobin(*iterables):
     iterables is small).
 
     """
-    # Recipe credited to George Sakkis
-    pending = len(iterables)
-    nexts = cycle(iter(it).__next__ for it in iterables)
-    while pending:
-        try:
-            for next in nexts:
-                yield next()
-        except StopIteration:
-            pending -= 1
-            nexts = cycle(islice(nexts, pending))
+    # Algorithm credited to George Sakkis
+    iterators = map(iter, iterables)
+    for num_active in range(len(iterables), 0, -1):
+        iterators = cycle(islice(iterators, num_active))
+        yield from map(next, iterators)
 
 
 def partition(pred, iterable):
@@ -459,16 +461,14 @@ def powerset(iterable):
 
     :func:`powerset` will operate on iterables that aren't :class:`set`
     instances, so repeated elements in the input will produce repeated elements
-    in the output. Use :func:`unique_everseen` on the input to avoid generating
-    duplicates:
+    in the output.
 
         >>> seq = [1, 1, 0]
         >>> list(powerset(seq))
         [(), (1,), (1,), (0,), (1, 1), (1, 0), (1, 0), (1, 1, 0)]
-        >>> from more_itertools import unique_everseen
-        >>> list(powerset(unique_everseen(seq)))
-        [(), (1,), (0,), (1, 0)]
 
+    For a variant that efficiently yields actual :class:`set` instances, see
+    :func:`powerset_of_sets`.
     """
     s = list(iterable)
     return chain.from_iterable(combinations(s, r) for r in range(len(s) + 1))
@@ -828,8 +828,6 @@ def iter_index(iterable, value, start=0, stop=None):
     """Yield the index of each place in *iterable* that *value* occurs,
     beginning with index *start* and ending before index *stop*.
 
-    See :func:`locate` for a more general means of finding the indexes
-    associated with particular values.
 
     >>> list(iter_index('AABCADEAF', 'A'))
     [0, 1, 4, 7]
@@ -837,6 +835,19 @@ def iter_index(iterable, value, start=0, stop=None):
     [1, 4, 7]
     >>> list(iter_index('AABCADEAF', 'A', 1, 7))  # stop index is not inclusive
     [1, 4]
+
+    The behavior for non-scalar *values* matches the built-in Python types.
+
+    >>> list(iter_index('ABCDABCD', 'AB'))
+    [0, 4]
+    >>> list(iter_index([0, 1, 2, 3, 0, 1, 2, 3], [0, 1]))
+    []
+    >>> list(iter_index([[0, 1], [2, 3], [0, 1], [2, 3]], [0, 1]))
+    [0, 2]
+
+    See :func:`locate` for a more general means of finding the indexes
+    associated with particular values.
+
     """
     seq_index = getattr(iterable, 'index', None)
     if seq_index is None:
@@ -1007,7 +1018,9 @@ def totient(n):
     >>> totient(12)
     4
     """
-    for p in unique_justseen(factor(n)):
+    # The itertools docs use unique_justseen instead of set; see
+    # https://github.com/more-itertools/more-itertools/issues/823
+    for p in set(factor(n)):
         n = n // p * (p - 1)
 
     return n
diff --git a/more_itertools/recipes.pyi b/more_itertools/recipes.pyi
index 15c6777..5307ecf 100644
--- a/more_itertools/recipes.pyi
+++ b/more_itertools/recipes.pyi
@@ -29,7 +29,9 @@ def consume(iterator: Iterable[_T], n: int | None = ...) -> None: ...
 def nth(iterable: Iterable[_T], n: int) -> _T | None: ...
 @overload
 def nth(iterable: Iterable[_T], n: int, default: _U) -> _T | _U: ...
-def all_equal(iterable: Iterable[_T]) -> bool: ...
+def all_equal(
+    iterable: Iterable[_T], key: Callable[[_T], _U] | None = ...
+) -> bool: ...
 def quantify(
     iterable: Iterable[_T], pred: Callable[[_T], bool] = ...
 ) -> int: ...


Test patch:
diff --git a/tests/test_more.py b/tests/test_more.py
index 56a8637..0960190 100644
--- a/tests/test_more.py
+++ b/tests/test_more.py
@@ -5761,3 +5761,24 @@ class FilterMapTests(TestCase):
         )
         expected = [1, 2, 3]
         self.assertEqual(actual, expected)
+
+
+class PowersetOfSetsTests(TestCase):
+    def test_simple(self):
+        iterable = [0, 1, 2]
+        actual = list(mi.powerset_of_sets(iterable))
+        expected = [set(), {0}, {1}, {2}, {0, 1}, {0, 2}, {1, 2}, {0, 1, 2}]
+        self.assertEqual(actual, expected)
+
+    def test_hash_count(self):
+        hash_count = 0
+
+        class Str(str):
+            def __hash__(true_self):
+                nonlocal hash_count
+                hash_count += 1
+                return super.__hash__(true_self)
+
+        iterable = map(Str, 'ABBBCDD')
+        self.assertEqual(len(list(mi.powerset_of_sets(iterable))), 128)
+        self.assertLessEqual(hash_count, 14)
diff --git a/tests/test_recipes.py b/tests/test_recipes.py
index 92a5722..a127703 100644
--- a/tests/test_recipes.py
+++ b/tests/test_recipes.py
@@ -134,33 +134,30 @@ class NthTests(TestCase):
 
 
 class AllEqualTests(TestCase):
-    """Tests for ``all_equal()``"""
-
     def test_true(self):
-        """Everything is equal"""
         self.assertTrue(mi.all_equal('aaaaaa'))
         self.assertTrue(mi.all_equal([0, 0, 0, 0]))
 
     def test_false(self):
-        """Not everything is equal"""
         self.assertFalse(mi.all_equal('aaaaab'))
         self.assertFalse(mi.all_equal([0, 0, 0, 1]))
 
     def test_tricky(self):
-        """Not everything is identical, but everything is equal"""
         items = [1, complex(1, 0), 1.0]
         self.assertTrue(mi.all_equal(items))
 
     def test_empty(self):
-        """Return True if the iterable is empty"""
         self.assertTrue(mi.all_equal(''))
         self.assertTrue(mi.all_equal([]))
 
     def test_one(self):
-        """Return True if the iterable is singular"""
         self.assertTrue(mi.all_equal('0'))
         self.assertTrue(mi.all_equal([0]))
 
+    def test_key(self):
+        self.assertTrue(mi.all_equal('4٤໔４৪', key=int))
+        self.assertFalse(mi.all_equal('Abc', key=str.casefold))
+
 
 class QuantifyTests(TestCase):
     """Tests for ``quantify()``"""