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

Issue:
allow passing your own key function to sort_together
I recently came across a situation in which I wanted to use `sort_together`, but I also wanted to sort by a custom key.
What do you think about exposing a `key` argument to `sort_together`?

I think it would also be nice for the user's key function to take only the parameters corresponding to the `key_list`, so the common case of sorting by the first is the easiest, and so the two features of `key` and `key_list` work together.

Something to this effect:

```python
def sort_together(iterables, key_list=(0,), key=lambda *x: x, reverse=False):
    return list(
        zip(
            *sorted(
                zip(*iterables), key=lambda x: key(*always_iterable(itemgetter(*key_list)(x))), reverse=reverse
            )
        )
    )


list_of_x = [1, 2, 3]
list_of_y = [4, 5, 6]

sort_together((list_of_x , list_of_y), key=lambda x: x*2)
sort_together((list_of_x , list_of_y), key_list=(1,), key=lambda y: y*2)
sort_together((list_of_x , list_of_y), key_list=(0, 1), key=lambda x, y: x*y)
```

Happy to open a PR if you like the idea.

Pull request:
add key argument to sort_together() function
Closes #480

Patch:
diff --git a/more_itertools/more.py b/more_itertools/more.py
index 158f3b3..2db826d 100755
--- a/more_itertools/more.py
+++ b/more_itertools/more.py
@@ -1569,7 +1569,7 @@ def zip_offset(*iterables, offsets, longest=False, fillvalue=None):
     return zip(*staggered)
 
 
-def sort_together(iterables, key_list=(0,), reverse=False):
+def sort_together(iterables, key_list=(0,), key=None, reverse=False):
     """Return the input iterables sorted together, with *key_list* as the
     priority for sorting. All iterables are trimmed to the length of the
     shortest one.
@@ -1591,18 +1591,47 @@ def sort_together(iterables, key_list=(0,), reverse=False):
         >>> sort_together(iterables, key_list=(1, 2))
         [(2, 3, 1), (0, 0, 1), ('a', 'c', 'b')]
 
+    To sort by a function of the elements of the iterable, pass a *key*
+    function. Its arguments are the elements of the iterables corresponding to
+    the key list::
+
+        >>> names = ('a', 'b', 'c')
+        >>> lengths = (1, 2, 3)
+        >>> widths = (5, 2, 1)
+        >>> def area(length, width):
+        ...     return length * width
+        >>> sort_together([names, lengths, widths], key_list=(1, 2), key=area)
+        [('c', 'b', 'a'), (3, 2, 1), (1, 2, 5)]
+
     Set *reverse* to ``True`` to sort in descending order.
 
         >>> sort_together([(1, 2, 3), ('c', 'b', 'a')], reverse=True)
         [(3, 2, 1), ('a', 'b', 'c')]
 
     """
-    return list(
-        zip(
-            *sorted(
-                zip(*iterables), key=itemgetter(*key_list), reverse=reverse
+    if key is None:
+        # if there is no key function, the key argument to sorted is an
+        # itemgetter
+        key_argument = itemgetter(*key_list)
+    else:
+        # if there is a key function, call it with the items at the offsets
+        # specified by the key function as arguments
+        key_list = list(key_list)
+        if len(key_list) == 1:
+            # if key_list contains a single item, pass the item at that offset
+            # as the only argument to the key function
+            key_offset = key_list[0]
+            key_argument = lambda zipped_items: key(zipped_items[key_offset])
+        else:
+            # if key_list contains multiple items, use itemgetter to return a
+            # tuple of items, which we pass as *args to the key function
+            get_key_items = itemgetter(*key_list)
+            key_argument = lambda zipped_items: key(
+                *get_key_items(zipped_items)
             )
-        )
+
+    return list(
+        zip(*sorted(zip(*iterables), key=key_argument, reverse=reverse))
     )
 
 
@@ -2559,6 +2588,27 @@ class seekable:
         >>> next(it), next(it), next(it)
         ('0', '1', '2')
 
+    Call :meth:`peek` to look ahead one item without advancing the iterator:
+
+        >>> it = seekable('1234')
+        >>> it.peek()
+        '1'
+        >>> list(it)
+        ['1', '2', '3', '4']
+        >>> it.peek(default='empty')
+        'empty'
+
+    Before the iterator is at its end, calling :func:`bool` on it will return
+    ``True``. After it will return ``False``:
+
+        >>> it = seekable('5678')
+        >>> bool(it)
+        True
+        >>> list(it)
+        ['5', '6', '7', '8']
+        >>> bool(it)
+        False
+
     You may view the contents of the cache with the :meth:`elements` method.
     That returns a :class:`SequenceView`, a view that updates automatically:
 
@@ -2616,6 +2666,25 @@ class seekable:
         self._cache.append(item)
         return item
 
+    def __bool__(self):
+        try:
+            self.peek()
+        except StopIteration:
+            return False
+        return True
+
+    def peek(self, default=_marker):
+        try:
+            peeked = next(self)
+        except StopIteration:
+            if default is _marker:
+                raise
+            return default
+        if self._index is None:
+            self._index = len(self._cache)
+        self._index -= 1
+        return peeked
+
     def elements(self):
         return SequenceView(self._cache)
 
@@ -2995,9 +3064,11 @@ def set_partitions(iterable, k=None):
         yield from set_partitions_helper(L, k)
 
 
-def time_limited(limit_seconds, iterable):
+class time_limited:
     """
     Yield items from *iterable* until *limit_seconds* have passed.
+    If the time limit expires before all items have been yielded, the
+    ``timed_out`` parameter will be set to ``True``.
 
     >>> from time import sleep
     >>> def generator():
@@ -3005,9 +3076,11 @@ def time_limited(limit_seconds, iterable):
     ...     yield 2
     ...     sleep(0.2)
     ...     yield 3
-    >>> iterable = generator()
-    >>> list(time_limited(0.1, iterable))
+    >>> iterable = time_limited(0.1, generator())
+    >>> list(iterable)
     [1, 2]
+    >>> iterable.timed_out
+    True
 
     Note that the time is checked before each item is yielded, and iteration
     stops if  the time elapsed is greater than *limit_seconds*. If your time
@@ -3015,14 +3088,25 @@ def time_limited(limit_seconds, iterable):
     the iterable, the function will run for 2 seconds and not yield anything.
 
     """
-    if limit_seconds < 0:
-        raise ValueError('limit_seconds must be positive')
 
-    start_time = monotonic()
-    for item in iterable:
-        if monotonic() - start_time > limit_seconds:
-            break
-        yield item
+    def __init__(self, limit_seconds, iterable):
+        if limit_seconds < 0:
+            raise ValueError('limit_seconds must be positive')
+        self.limit_seconds = limit_seconds
+        self._iterable = iter(iterable)
+        self._start_time = monotonic()
+        self.timed_out = False
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        item = next(self._iterable)
+        if monotonic() - self._start_time > self.limit_seconds:
+            self.timed_out = True
+            raise StopIteration
+
+        return item
 
 
 def only(iterable, default=None, too_long=None):
diff --git a/more_itertools/more.pyi b/more_itertools/more.pyi
index e0c3d95..f16feca 100644
--- a/more_itertools/more.pyi
+++ b/more_itertools/more.pyi
@@ -216,6 +216,7 @@ def zip_offset(
 def sort_together(
     iterables: Iterable[Iterable[_T]],
     key_list: Iterable[int] = ...,
+    key: Optional[Callable[..., Any]] = ...,
     reverse: bool = ...,
 ) -> List[Tuple[_T, ...]]: ...
 def unzip(iterable: Iterable[Sequence[_T]]) -> Tuple[Iterator[_T], ...]: ...
@@ -322,6 +323,11 @@ class seekable(Generic[_T], Iterator[_T]):
     ) -> None: ...
     def __iter__(self) -> seekable[_T]: ...
     def __next__(self) -> _T: ...
+    def __bool__(self) -> bool: ...
+    @overload
+    def peek(self) -> _T: ...
+    @overload
+    def peek(self, default: _U) -> Union[_T, _U]: ...
     def elements(self) -> SequenceView[_T]: ...
     def seek(self, index: int) -> None: ...
 
@@ -382,9 +388,14 @@ def partitions(iterable: Iterable[_T]) -> Iterator[List[List[_T]]]: ...
 def set_partitions(
     iterable: Iterable[_T], k: Optional[int] = ...
 ) -> Iterator[List[List[_T]]]: ...
-def time_limited(
-    limit_seconds: float, iterable: Iterable[_T]
-) -> Iterator[_T]: ...
+
+class time_limited(Generic[_T], Iterator[_T]):
+    def __init__(
+        self, limit_seconds: float, iterable: Iterable[_T]
+    ) -> None: ...
+    def __iter__(self) -> islice_extended[_T]: ...
+    def __next__(self) -> _T: ...
+
 @overload
 def only(
     iterable: Iterable[_T], *, too_long: Optional[_Raisable] = ...
diff --git a/more_itertools/recipes.py b/more_itertools/recipes.py
index da57644..c06f47b 100644
--- a/more_itertools/recipes.py
+++ b/more_itertools/recipes.py
@@ -386,12 +386,13 @@ def unique_everseen(iterable, key=None):
     ``key=lambda x: frozenset(x.items())`` can be used.
 
     """
+    key = key if key is not None else lambda x: x
     seenset = set()
     seenset_add = seenset.add
     seenlist = []
     seenlist_add = seenlist.append
-    iterable, keys = tee(iterable)
-    for element, k in zip(iterable, map(key, keys) if key else keys):
+    for element in iterable:
+        k = key(element)
         try:
             if k not in seenset:
                 seenset_add(k)


Test patch:
diff --git a/tests/test_more.py b/tests/test_more.py
index 4885826..d34d884 100644
--- a/tests/test_more.py
+++ b/tests/test_more.py
@@ -223,15 +223,22 @@ class NthOrLastTests(TestCase):
         self.assertRaises(ValueError, lambda: mi.nth_or_last(range(0), 0))
 
 
-class PeekableTests(TestCase):
-    """Tests for ``peekable()`` behavor not incidentally covered by testing
-    ``collate()``
+class PeekableMixinTests:
+    """Common tests for ``peekable()`` and ``seekable()`` behavior"""
 
-    """
+    cls = None
+
+    def test_passthrough(self):
+        """Iterating a peekable without using ``peek()`` or ``prepend()``
+        should just give the underlying iterable's elements (a trivial test but
+        useful to set a baseline in case something goes wrong)"""
+        expected = [1, 2, 3, 4, 5]
+        actual = list(self.cls(expected))
+        self.assertEqual(actual, expected)
 
     def test_peek_default(self):
         """Make sure passing a default into ``peek()`` works."""
-        p = mi.peekable([])
+        p = self.cls([])
         self.assertEqual(p.peek(7), 7)
 
     def test_truthiness(self):
@@ -239,10 +246,10 @@ class PeekableTests(TestCase):
         the iterable.
 
         """
-        p = mi.peekable([])
+        p = self.cls([])
         self.assertFalse(p)
 
-        p = mi.peekable(range(3))
+        p = self.cls(range(3))
         self.assertTrue(p)
 
     def test_simple_peeking(self):
@@ -250,11 +257,21 @@ class PeekableTests(TestCase):
         iterator, respectively.
 
         """
-        p = mi.peekable(range(10))
+        p = self.cls(range(10))
         self.assertEqual(next(p), 0)
         self.assertEqual(p.peek(), 1)
+        self.assertEqual(p.peek(), 1)
         self.assertEqual(next(p), 1)
 
+
+class PeekableTests(PeekableMixinTests, TestCase):
+    """Tests for ``peekable()`` behavior not incidentally covered by testing
+    ``collate()``
+
+    """
+
+    cls = mi.peekable
+
     def test_indexing(self):
         """
         Indexing into the peekable shouldn't advance the iterator.
@@ -334,14 +351,6 @@ class PeekableTests(TestCase):
         self.assertEqual(old_cache, list(p._cache))
         self.assertEqual(list(p), list(iterable))
 
-    def test_passthrough(self):
-        """Iterating a peekable without using ``peek()`` or ``prepend()``
-        should just give the underlying iterable's elements (a trivial test but
-        useful to set a baseline in case something goes wrong)"""
-        expected = [1, 2, 3, 4, 5]
-        actual = list(mi.peekable(expected))
-        self.assertEqual(actual, expected)
-
     # prepend() behavior tests
 
     def test_prepend(self):
@@ -1920,6 +1929,44 @@ class SortTogetherTest(TestCase):
             IndexError, lambda: mi.sort_together(iterables, key_list=(5,))
         )
 
+    def test_key_function(self):
+        """tests `key` function, including interaction with `key_list`"""
+        iterables = [
+            ['GA', 'GA', 'GA', 'CT', 'CT', 'CT'],
+            ['May', 'Aug.', 'May', 'June', 'July', 'July'],
+            [97, 20, 100, 70, 100, 20],
+        ]
+        self.assertEqual(
+            mi.sort_together(iterables, key=lambda x: x),
+            [
+                ('CT', 'CT', 'CT', 'GA', 'GA', 'GA'),
+                ('June', 'July', 'July', 'May', 'Aug.', 'May'),
+                (70, 100, 20, 97, 20, 100),
+            ],
+        )
+        self.assertEqual(
+            mi.sort_together(iterables, key=lambda x: x[::-1]),
+            [
+                ('GA', 'GA', 'GA', 'CT', 'CT', 'CT'),
+                ('May', 'Aug.', 'May', 'June', 'July', 'July'),
+                (97, 20, 100, 70, 100, 20),
+            ],
+        )
+        self.assertEqual(
+            mi.sort_together(
+                iterables,
+                key_list=(0, 2),
+                key=lambda state, number: number
+                if state == 'CT'
+                else 2 * number,
+            ),
+            [
+                ('CT', 'GA', 'CT', 'CT', 'GA', 'GA'),
+                ('July', 'Aug.', 'June', 'July', 'May', 'May'),
+                (20, 20, 70, 100, 97, 100),
+            ],
+        )
+
     def test_reverse(self):
         """tests `reverse` to ensure a reverse sort for `key_list` iterables"""
         iterables = [
@@ -3045,7 +3092,9 @@ class DifferenceTest(TestCase):
         self.assertEqual(actual, original)
 
 
-class SeekableTest(TestCase):
+class SeekableTest(PeekableMixinTests, TestCase):
+    cls = mi.seekable
+
     def test_exhaustion_reset(self):
         iterable = [str(n) for n in range(10)]
 
@@ -3663,16 +3712,25 @@ class TimeLimitedTests(TestCase):
             sleep(0.2)
             yield 3
 
-        iterable = generator()
-        actual = list(mi.time_limited(0.1, iterable))
+        iterable = mi.time_limited(0.1, generator())
+        actual = list(iterable)
         expected = [1, 2]
         self.assertEqual(actual, expected)
+        self.assertTrue(iterable.timed_out)
+
+    def test_complete(self):
+        iterable = mi.time_limited(2, iter(range(10)))
+        actual = list(iterable)
+        expected = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+        self.assertEqual(actual, expected)
+        self.assertFalse(iterable.timed_out)
 
     def test_zero_limit(self):
-        iterable = count()
-        actual = list(mi.time_limited(0, iterable))
+        iterable = mi.time_limited(0, count())
+        actual = list(iterable)
         expected = []
         self.assertEqual(actual, expected)
+        self.assertTrue(iterable.timed_out)
 
     def test_invalid_limit(self):
         with self.assertRaises(ValueError):