diff --git a/docs/conf.py b/docs/conf.py index 63766e0..f982072 100644 --- a/docs/conf.py +++ b/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'), +] diff --git a/jaraco/collections/__init__.py b/jaraco/collections/__init__.py index 2b9f62a..0d501cf 100644 --- a/jaraco/collections/__init__.py +++ b/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,14 +265,17 @@ 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]) @@ -254,7 +283,7 @@ def bounds(self): undefined_value = type('RangeValueUndefined', (), {})() class Item(int): - "RangeMap Item" + """RangeMap Item""" first_item = Item(0) last_item = Item(-1) diff --git a/newsfragments/16.feature.rst b/newsfragments/16.feature.rst new file mode 100644 index 0000000..740b18c --- /dev/null +++ b/newsfragments/16.feature.rst @@ -0,0 +1 @@ +Fully typed ``RangeMap`` and avoid complete iterations to find matches -- by :user:`Avasam`