add rank(), rank_by()

morelinq/MoreLINQ#826

doc/api/more/types_linq.more.more_enumerable.rst

@@ -334,6 +334,86 @@ Example
 
 ----
 
+instancemethod ``rank[TSupportsLessThan]()``
+----------------------------------------------
+
+Constraint
+  - `self`: ``MoreEnumerable[TSupportsLessThan]``
+
+Returns
+  - ``MoreEnumerable[int]``
+
+Ranks each item in the sequence in descending order.
+
+Example
+    >>> scores = [1, 4, 77, 23, 23, 4, 9, 0, -7, 101, 23]
+    >>> MoreEnumerable(scores).rank().to_list()
+    [6, 5, 2, 3, 3, 5, 4, 7, 8, 1, 3]  # 101 is largest, so has rank of 1
+
+----
+
+instancemethod ``rank(__comparer)``
+-------------------------------------
+
+Parameters
+  - `__comparer` (``Callable[[TSource_co, TSource_co], int]``)
+
+Returns
+  - ``MoreEnumerable[int]``
+
+Ranks each item in the sequence in descending order using the given comparer.
+
+Such comparer takes two values and return positive ints when lhs > rhs, negative ints
+if lhs < rhs, and 0 if they are equal.
+
+----
+
+instancemethod ``rank_by[TSupportsLessThan](key_selector)``
+-------------------------------------------------------------
+
+Parameters
+  - `key_selector` (``Callable[[TSource_co], TSupportsLessThan]``)
+
+Returns
+  - ``MoreEnumerable[int]``
+
+Ranks each item in the sequence in descending order using the given selector.
+
+Example
+    .. code-block:: python
+
+        >>> scores = [
+        ...     {'name': 'Frank', 'score': 75},
+        ...     {'name': 'Alica', 'score': 90},
+        ...     {'name': 'Erika', 'score': 99},
+        ...     {'name': 'Rogers', 'score': 90},
+        ... ]
+
+        >>> MoreEnumerable(scores).rank_by(lambda x: x['score']) \
+        ...     .zip(scores) \
+        ...     .group_by(lambda t: t[0], lambda t: t[1]['name']) \
+        ...     .to_dict(lambda g: g.key, lambda g: g.to_list())
+        {3: ['Frank'], 2: ['Alica', 'Rogers'], 1: ['Erika']}
+
+----
+
+instancemethod ``rank_by[TKey](key_selector, __comparer)``
+------------------------------------------------------------
+
+Parameters
+  - `key_selector` (``Callable[[TSource_co], TKey]``)
+  - `__comparer` (``Callable[[TKey, TKey], int]``)
+
+Returns
+  - ``MoreEnumerable[int]``
+
+Ranks each item in the sequence in descending order using the given selector and comparer.
+
+Such comparer takes two values and return positive ints when lhs > rhs, negative ints
+if lhs < rhs, and 0 if they are equal.
+
+----
+
 staticmethod ``traverse_breath_first[TSource](root, children_selector)``
 --------------------------------------------------------------------------
 

doc/api_spec.py

@@ -194,6 +194,8 @@ class ClassSpec(TypedDict):
                     'maxima_by',
                     'minima_by',
                     'pipe',
+                    'rank',
+                    'rank_by',
                     'traverse_breath_first',
                     'traverse_depth_first',
                 },

tests/test_more_usage.py

@@ -373,6 +373,51 @@ def test_action_then_yield(self):
         assert q.to_list() == []
 
 
+class TestRankMethod:
+    def test_overload1(self):
+        def gen():
+            yield from [1, 4, 77, 23, 23, 4, 9, 0, -7, 101, 23]
+        en = MoreEnumerable(gen())
+        assert en.rank().to_list() == [6, 5, 2, 3, 3, 5, 4, 7, 8, 1, 3]
+
+    def test_overload1_empty(self):
+        assert MoreEnumerable([]).rank().to_list() == []
+
+    def test_overload1_sorted(self):
+        ints = [444, 190, 129, 122, 100]
+        assert MoreEnumerable(ints).rank().to_list() == [1, 2, 3, 4, 5]
+        assert MoreEnumerable(reversed(ints)).rank().to_list() == [5, 4, 3, 2, 1]
+
+    def test_overload1_same(self):
+        en = MoreEnumerable([8, 8, 8])
+        assert en.rank().to_list() == [1, 1, 1]
+
+    def test_overload2(self):
+        en = MoreEnumerable([(1, ''), (1, ''), (4, ''), (4, ''), (3, '')])
+        assert en.rank(lambda lhs, rhs: lhs[0] - rhs[0]).to_list() == [3, 3, 1, 1, 2]
+
+
+# majority is already tested by rank test
+class TestRankByMethod:
+    def test_overload1(self):
+        def gen():
+            yield from ['aaa', 'xyz', 'carbon', 'emission', 'statistics', 'somany']
+        en = MoreEnumerable(gen())
+        assert en.rank_by(len).to_list() == [4, 4, 3, 2, 1, 3]
+
+    def test_overload2(self):
+        en = MoreEnumerable([
+            ['aaa'], ['xyz'], ['carbon'], ['emission'], ['statistics'], ['somany']
+        ])
+        assert en.rank_by(lambda x: x[0], lambda lhs, rhs: len(lhs) - len(rhs)) \
+            .to_list() == [4, 4, 3, 2, 1, 3]
+
+    def test_overload2_empty(self):
+        en = MoreEnumerable([])
+        assert en.rank_by(lambda x: x[0], lambda lhs, rhs: len(lhs) - len(rhs)) \
+            .to_list() == []
+
+
 class TestTraverseBreathFirstMethod:
     tree = Tree \
     (

types_linq/more/more_enumerable.py

@@ -6,7 +6,7 @@
     from .extrema_enumerable import ExtremaEnumerable
 
 from ..enumerable import Enumerable
-from ..util import ComposeSet
+from ..util import ComposeMap, ComposeSet
 from ..more_typing import (
     TKey,
     TSource,
@@ -136,6 +136,41 @@ def inner():
                 yield elem
         return MoreEnumerable(inner)
 
+    def rank(self, *args: Callable[[TSource_co, TSource_co], int]) -> MoreEnumerable[int]:
+        return self.rank_by(lambda x: x, *args)
+
+    def rank_by(self,
+        key_selector: Callable[[TSource_co], TKey],
+        *args: Callable[[TKey, TKey], int]) -> MoreEnumerable[int]:
+        if len(args) == 0:
+            comparer = None
+        else:  # len(args) == 1
+            comparer = args[0]
+
+        def inner():
+            # avoid enumerating twice
+            copy = MoreEnumerable(self.select(key_selector).to_list())
+            ordered = copy.distinct() \
+                .order_by_descending(lambda x: x, *args)
+            if comparer is None:
+                # replaces .enumerate()
+                rank_map = ComposeMap(ordered.select2(lambda x, i: (x, i + 1)))
+            else:
+                # this is different from morelinq
+                ordered = ordered.to_list()
+                if not ordered:
+                    return
+                rank_map = ComposeMap()
+                rank_map[ordered[0]] = 1
+                rank = 1
+                for i in range(1, len(ordered)):
+                    if comparer(ordered[i - 1], ordered[i]) != 0:
+                        rank += 1
+                    rank_map[ordered[i]] = rank
+            for key in copy:
+                yield rank_map[key]
+        return MoreEnumerable(inner)
+
     @staticmethod
     def traverse_breath_first(
         root: TSource,

types_linq/more/more_enumerable.pyi

@@ -224,6 +224,60 @@ class MoreEnumerable(Enumerable[TSource_co]):
             {1, 2}
         '''
 
+    @overload
+    def rank(self: MoreEnumerable[TSupportsLessThan]) -> MoreEnumerable[int]:
+        '''
+        Ranks each item in the sequence in descending order.
+
+        Example
+            >>> scores = [1, 4, 77, 23, 23, 4, 9, 0, -7, 101, 23]
+            >>> MoreEnumerable(scores).rank().to_list()
+            [6, 5, 2, 3, 3, 5, 4, 7, 8, 1, 3]  # 101 is largest, so has rank of 1
+        '''
+
+    @overload
+    def rank(self, __comparer: Callable[[TSource_co, TSource_co], int]) -> MoreEnumerable[int]:
+        '''
+        Ranks each item in the sequence in descending order using the given comparer.
+
+        Such comparer takes two values and return positive ints when lhs > rhs, negative ints
+        if lhs < rhs, and 0 if they are equal.
+        '''
+
+    @overload
+    def rank_by(self, key_selector: Callable[[TSource_co], TSupportsLessThan]) -> MoreEnumerable[int]:
+        '''
+        Ranks each item in the sequence in descending order using the given selector.
+
+        Example
+            .. code-block:: python
+
+                >>> scores = [
+                ...     {'name': 'Frank', 'score': 75},
+                ...     {'name': 'Alica', 'score': 90},
+                ...     {'name': 'Erika', 'score': 99},
+                ...     {'name': 'Rogers', 'score': 90},
+                ... ]
+
+                >>> MoreEnumerable(scores).rank_by(lambda x: x['score']) \\
+                ...     .zip(scores) \\
+                ...     .group_by(lambda t: t[0], lambda t: t[1]['name']) \\
+                ...     .to_dict(lambda g: g.key, lambda g: g.to_list())
+                {3: ['Frank'], 2: ['Alica', 'Rogers'], 1: ['Erika']}
+        '''
+
+    @overload
+    def rank_by(self,
+        key_selector: Callable[[TSource_co], TKey],
+        __comparer: Callable[[TKey, TKey], int],
+    ) -> MoreEnumerable[int]:
+        '''
+        Ranks each item in the sequence in descending order using the given selector and comparer.
+
+        Such comparer takes two values and return positive ints when lhs > rhs, negative ints
+        if lhs < rhs, and 0 if they are equal.
+        '''
+
     @staticmethod
     def traverse_breath_first(
         root: TSource,