Fully typed RangeMap and avoid complete iterations to find matches (#…

…16)

Co-authored-by: Jason R. Coombs <jaraco@jaraco.com>

docs/conf.py

@@ -54,3 +54,16 @@
 nitpick_ignore += [
     ('py:class', 're.Pattern'),
 ]
+
+# jaraco/jaraco.collections#16
+nitpick_ignore += [
+    ('py:class', 'SupportsKeysAndGetItem'),
+    ('py:class', '_RangeMapKT'),
+    ('py:class', '_VT'),
+    ('py:class', '_T'),
+    ('py:class', 'jaraco.collections._RangeMapKT'),
+    ('py:class', 'jaraco.collections._VT'),
+    ('py:class', 'jaraco.collections._T'),
+    ('py:obj', 'jaraco.collections._RangeMapKT'),
+    ('py:obj', 'jaraco.collections._VT'),
+]

jaraco/collections/__init__.py

@@ -8,10 +8,25 @@
 import random
 import re
 from collections.abc import Container, Iterable, Mapping
-from typing import Any, Callable, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, TypeVar, Union, overload
 
 import jaraco.text
 
+if TYPE_CHECKING:
+    from _operator import _SupportsComparison
+
+    from _typeshed import SupportsKeysAndGetItem
+    from typing_extensions import Self
+
+    _RangeMapKT = TypeVar('_RangeMapKT', bound=_SupportsComparison)
+else:
+    # _SupportsComparison doesn't exist at runtime,
+    # but _RangeMapKT is used in RangeMap's superclass' type parameters
+    _RangeMapKT = TypeVar('_RangeMapKT')
+
+_T = TypeVar('_T')
+_VT = TypeVar('_VT')
+
 _Matchable = Union[Callable, Container, Iterable, re.Pattern]
 
 
@@ -120,7 +135,7 @@ def dict_map(function, dictionary):
     return dict((key, function(value)) for key, value in dictionary.items())
 
 
-class RangeMap(dict):
+class RangeMap(Dict[_RangeMapKT, _VT]):
     """
     A dictionary-like object that uses the keys as bounds for a range.
     Inclusion of the value for that range is determined by the
@@ -203,21 +218,28 @@ class RangeMap(dict):
 
     def __init__(
         self,
-        source,
+        source: (
+            SupportsKeysAndGetItem[_RangeMapKT, _VT] | Iterable[tuple[_RangeMapKT, _VT]]
+        ),
         sort_params: Mapping[str, Any] = {},
-        key_match_comparator=operator.le,
+        key_match_comparator: Callable[[_RangeMapKT, _RangeMapKT], bool] = operator.le,
     ):
         dict.__init__(self, source)
         self.sort_params = sort_params
         self.match = key_match_comparator
 
     @classmethod
-    def left(cls, source):
+    def left(
+        cls,
+        source: (
+            SupportsKeysAndGetItem[_RangeMapKT, _VT] | Iterable[tuple[_RangeMapKT, _VT]]
+        ),
+    ) -> Self:
         return cls(
             source, sort_params=dict(reverse=True), key_match_comparator=operator.ge
         )
 
-    def __getitem__(self, item):
+    def __getitem__(self, item: _RangeMapKT) -> _VT:
         sorted_keys = sorted(self.keys(), **self.sort_params)
         if isinstance(item, RangeMap.Item):
             result = self.__getitem__(sorted_keys[item])
@@ -228,7 +250,11 @@ def __getitem__(self, item):
                 raise KeyError(key)
         return result
 
-    def get(self, key, default=None):
+    @overload  # type: ignore[override] # Signature simplified over dict and Mapping
+    def get(self, key: _RangeMapKT, default: _T) -> _VT | _T: ...
+    @overload
+    def get(self, key: _RangeMapKT, default: None = None) -> _VT | None: ...
+    def get(self, key: _RangeMapKT, default: _T | None = None) -> _VT | _T | None:
         """
         Return the value for key if key is in the dictionary, else default.
         If default is not given, it defaults to None, so that this method
@@ -239,22 +265,25 @@ def get(self, key, default=None):
         except KeyError:
             return default
 
-    def _find_first_match_(self, keys, item):
+    def _find_first_match_(
+        self, keys: Iterable[_RangeMapKT], item: _RangeMapKT
+    ) -> _RangeMapKT:
         is_match = functools.partial(self.match, item)
-        matches = list(filter(is_match, keys))
-        if matches:
-            return matches[0]
-        raise KeyError(item)
+        matches = filter(is_match, keys)
+        try:
+            return next(matches)
+        except StopIteration:
+            raise KeyError(item) from None
 
-    def bounds(self):
+    def bounds(self) -> tuple[_RangeMapKT, _RangeMapKT]:
         sorted_keys = sorted(self.keys(), **self.sort_params)
         return (sorted_keys[RangeMap.first_item], sorted_keys[RangeMap.last_item])
 
     # some special values for the RangeMap
     undefined_value = type('RangeValueUndefined', (), {})()
 
     class Item(int):
-        "RangeMap Item"
+        """RangeMap Item"""
 
     first_item = Item(0)
     last_item = Item(-1)

newsfragments/16.feature.rst

@@ -0,0 +1 @@
+Fully typed ``RangeMap`` and avoid complete iterations to find matches -- by :user:`Avasam`