Add type annotations to sortedcontainers

.gitignore

@@ -17,3 +17,7 @@
 
 # macOS metadata
 .DS_Store
+
+# typing cache
+.mypy_cache/
+.pyre/

setup.py

@@ -28,6 +28,7 @@ def run_tests(self):
     url='http://www.grantjenks.com/docs/sortedcontainers/',
     license='Apache 2.0',
     packages=['sortedcontainers'],
+    package_data={'sortedcontainers': ['py.typed', '*.pyi']},
     tests_require=['tox'],
     cmdclass={'test': Tox},
     install_requires=[],

sortedcontainers/sorteddict.pyi

@@ -0,0 +1,122 @@
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generic,
+    Hashable,
+    Iterator,
+    Iterable,
+    ItemsView,
+    KeysView,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Type,
+    TypeVar,
+    Tuple,
+    Union,
+    ValuesView,
+    overload,
+)
+
+_T = TypeVar("_T")
+_S = TypeVar("_S")
+_T_h = TypeVar("_T_h", bound=Hashable)
+_KT = TypeVar("_KT", bound=Hashable)  # Key type.
+_VT = TypeVar("_VT")  # Value type.
+_KT_co = TypeVar("_KT_co", covariant=True, bound=Hashable)
+_VT_co = TypeVar("_VT_co", covariant=True)
+_SD = TypeVar("_SD", bound=SortedDict)
+_Key = Callable[[_T], Any]
+
+class SortedDict(Dict[_KT, _VT]):
+    @overload
+    def __init__(self, **kwargs: _VT) -> None: ...
+    @overload
+    def __init__(self, __map: Mapping[_KT, _VT], **kwargs: _VT) -> None: ...
+    @overload
+    def __init__(
+        self, __iterable: Iterable[Tuple[_KT, _VT]], **kwargs: _VT
+    ) -> None: ...
+    @overload
+    def __init__(self, __key: _Key[_KT], **kwargs: _VT) -> None: ...
+    @overload
+    def __init__(
+        self, __key: _Key[_KT], __map: Mapping[_KT, _VT], **kwargs: _VT
+    ) -> None: ...
+    @overload
+    def __init__(
+        self,
+        __key: _Key[_KT],
+        __iterable: Iterable[Tuple[_KT, _VT]],
+        **kwargs: _VT
+    ) -> None: ...
+    @property
+    def key(self) -> Optional[_Key[_KT]]: ...
+    @property
+    def iloc(self) -> SortedKeysView[_KT]: ...
+    def clear(self) -> None: ...
+    def __delitem__(self, key: _KT) -> None: ...
+    def __iter__(self) -> Iterator[_KT]: ...
+    def __reversed__(self) -> Iterator[_KT]: ...
+    def __setitem__(self, key: _KT, value: _VT) -> None: ...
+    def _setitem(self, key: _KT, value: _VT) -> None: ...
+    def copy(self: _SD) -> _SD: ...
+    def __copy__(self: _SD) -> _SD: ...
+    @classmethod
+    @overload
+    def fromkeys(cls, seq: Iterable[_T_h]) -> SortedDict[_T_h, None]: ...
+    @classmethod
+    @overload
+    def fromkeys(cls, seq: Iterable[_T_h], value: _S) -> SortedDict[_T_h, _S]: ...
+    def keys(self) -> SortedKeysView[_KT]: ...
+    def items(self) -> SortedItemsView[_KT, _VT]: ...
+    def values(self) -> SortedValuesView[_VT]: ...
+    @overload
+    def pop(self, key: _KT) -> _VT: ...
+    @overload
+    def pop(self, key: _KT, default: _T = ...) -> Union[_VT, _T]: ...
+    def popitem(self, index: int = ...) -> Tuple[_KT, _VT]: ...
+    def peekitem(self, index: int = ...) -> Tuple[_KT, _VT]: ...
+    def setdefault(self, key: _KT, default: _VT = ...) -> _VT: ...
+    @overload
+    def update(self, __map: Mapping[_KT, _VT], **kwargs: _VT) -> None: ...
+    @overload
+    def update(
+        self, __iterable: Iterable[Tuple[_KT, _VT]], **kwargs: _VT
+    ) -> None: ...
+    @overload
+    def update(self, **kwargs: _VT) -> None: ...
+    def __reduce__(
+        self
+    ) -> Tuple[
+        Type[SortedDict[_KT, _VT]],
+        Tuple[Callable[[_KT], Any], List[Tuple[_KT, _VT]]],
+    ]: ...
+    def __repr__(self) -> str: ...
+    def _check(self) -> None: ...
+
+class SortedKeysView(KeysView[_KT_co], Sequence[_KT_co]):
+    @overload
+    def __getitem__(self, index: int) -> _KT_co: ...
+    @overload
+    def __getitem__(self, index: slice) -> List[_KT_co]: ...
+    def __delitem__(self, index: Union[int, slice]) -> None: ...
+
+class SortedItemsView(
+    ItemsView[_KT_co, _VT_co], Sequence[Tuple[_KT_co, _VT_co]]
+):
+    def __iter__(self) -> Iterator[Tuple[_KT_co, _VT_co]]: ...
+    @overload
+    def __getitem__(self, index: int) -> Tuple[_KT_co, _VT_co]: ...
+    @overload
+    def __getitem__(self, index: slice) -> List[Tuple[_KT_co, _VT_co]]: ...
+    def __delitem__(self, index: Union[int, slice]) -> None: ...
+
+class SortedValuesView(ValuesView[_VT_co], Sequence[_VT_co]):
+    @overload
+    def __getitem__(self, index: int) -> _VT_co: ...
+    @overload
+    def __getitem__(self, index: slice) -> List[_VT_co]: ...
+    def __delitem__(self, index: Union[int, slice]) -> None: ...

sortedcontainers/sortedlist.py

@@ -16,9 +16,6 @@
 # pylint: disable=too-many-lines
 from __future__ import print_function
 
-import sys
-import traceback
-
 from bisect import bisect_left, bisect_right, insort
 from itertools import chain, repeat, starmap
 from math import log
@@ -1670,6 +1667,8 @@ def _check(self):
                         child_sum = self._index[child] + self._index[child + 1]
                         assert child_sum == self._index[pos]
         except:
+            import sys # pylint: disable=import-outside-toplevel
+            import traceback # pylint: disable=import-outside-toplevel
             traceback.print_exc(file=sys.stdout)
             print('len', self._len)
             print('load', self._load)
@@ -2628,6 +2627,8 @@ def _check(self):
                         child_sum = self._index[child] + self._index[child + 1]
                         assert child_sum == self._index[pos]
         except:
+            import sys # pylint: disable=import-outside-toplevel
+            import traceback # pylint: disable=import-outside-toplevel
             traceback.print_exc(file=sys.stdout)
             print('len', self._len)
             print('load', self._load)

sortedcontainers/sortedlist.pyi

@@ -0,0 +1,187 @@
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Iterable,
+    Iterator,
+    List,
+    MutableSequence,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    overload,
+)
+
+_T = TypeVar("_T")
+_SL = TypeVar("_SL", bound=SortedList)
+_SKL = TypeVar("_SKL", bound=SortedKeyList)
+_Key = Callable[[_T], Any]
+_Repr = Callable[[], str]
+
+def recursive_repr(fillvalue: str = ...) -> Callable[[_Repr], _Repr]: ...
+
+class SortedList(MutableSequence[_T]):
+
+    DEFAULT_LOAD_FACTOR: int = ...
+    def __init__(
+        self,
+        iterable: Optional[Iterable[_T]] = ...,
+        key: Optional[_Key[_T]] = ...,
+    ): ...
+    # NB: currently mypy does not honour return type, see mypy #3307
+    @overload
+    def __new__(cls: Type[_SL], iterable: None, key: None) -> _SL: ...
+    @overload
+    def __new__(
+        cls: Type[_SL], iterable: None, key: _Key[_T]
+    ) -> SortedKeyList[_T]: ...
+    @overload
+    def __new__(cls: Type[_SL], iterable: Iterable[_T], key: None) -> _SL: ...
+    @overload
+    def __new__(
+        cls, iterable: Iterable[_T], key: _Key[_T]
+    ) -> SortedKeyList[_T]: ...
+    @property
+    def key(self) -> Optional[Callable[[_T], Any]]: ...
+    def _reset(self, load: int) -> None: ...
+    def clear(self) -> None: ...
+    def _clear(self) -> None: ...
+    def add(self, value: _T) -> None: ...
+    def _expand(self, pos: int) -> None: ...
+    def update(self, iterable: Iterable[_T]) -> None: ...
+    def _update(self, iterable: Iterable[_T]) -> None: ...
+    def discard(self, value: _T) -> None: ...
+    def remove(self, value: _T) -> None: ...
+    def _delete(self, pos: int, idx: int) -> None: ...
+    def _loc(self, pos: int, idx: int) -> int: ...
+    def _pos(self, idx: int) -> int: ...
+    def _build_index(self) -> None: ...
+    def __contains__(self, value: Any) -> bool: ...
+    def __delitem__(self, index: Union[int, slice]) -> None: ...
+    @overload
+    def __getitem__(self, index: int) -> _T: ...
+    @overload
+    def __getitem__(self, index: slice) -> List[_T]: ...
+    @overload
+    def _getitem(self, index: int) -> _T: ...
+    @overload
+    def _getitem(self, index: slice) -> List[_T]: ...
+    @overload
+    def __setitem__(self, index: int, value: _T) -> None: ...
+    @overload
+    def __setitem__(self, index: slice, value: Iterable[_T]) -> None: ...
+    def __iter__(self) -> Iterator[_T]: ...
+    def __reversed__(self) -> Iterator[_T]: ...
+    def __len__(self) -> int: ...
+    def reverse(self) -> None: ...
+    def islice(
+        self,
+        start: Optional[int] = ...,
+        stop: Optional[int] = ...,
+        reverse=bool,
+    ) -> Iterator[_T]: ...
+    def _islice(
+        self,
+        min_pos: int,
+        min_idx: int,
+        max_pos: int,
+        max_idx: int,
+        reverse: bool,
+    ) -> Iterator[_T]: ...
+    def irange(
+        self,
+        minimum: Optional[_T] = ...,
+        maximum: Optional[_T] = ...,
+        inclusive: Tuple[bool, bool] = ...,
+        reverse: bool = ...,
+    ) -> Iterator[_T]: ...
+    def bisect_left(self, value: _T) -> int: ...
+    def bisect_right(self, value: _T) -> int: ...
+    def bisect(self, value: _T) -> int: ...
+    def _bisect_right(self, value: _T) -> int: ...
+    def count(self, value: _T) -> int: ...
+    def copy(self: _SL) -> _SL: ...
+    def __copy__(self: _SL) -> _SL: ...
+    def append(self, value: _T) -> None: ...
+    def extend(self, values: Iterable[_T]) -> None: ...
+    def insert(self, index: int, value: _T) -> None: ...
+    def pop(self, index: int = ...) -> _T: ...
+    def index(
+        self, value: _T, start: Optional[int] = ..., stop: Optional[int] = ...
+    ) -> int: ...
+    def __add__(self: _SL, other: Iterable[_T]) -> _SL: ...
+    def __radd__(self: _SL, other: Iterable[_T]) -> _SL: ...
+    def __iadd__(self: _SL, other: Iterable[_T]) -> _SL: ...
+    def __mul__(self: _SL, num: int) -> _SL: ...
+    def __rmul__(self: _SL, num: int) -> _SL: ...
+    def __imul__(self: _SL, num: int) -> _SL: ...
+    def __eq__(self, other: Any) -> bool: ...
+    def __ne__(self, other: Any) -> bool: ...
+    def __lt__(self, other: Sequence[_T]) -> bool: ...
+    def __gt__(self, other: Sequence[_T]) -> bool: ...
+    def __le__(self, other: Sequence[_T]) -> bool: ...
+    def __ge__(self, other: Sequence[_T]) -> bool: ...
+    def __repr__(self) -> str: ...
+    def _check(self) -> None: ...
+
+class SortedKeyList(SortedList[_T]):
+    def __init__(
+        self, iterable: Optional[Iterable[_T]] = ..., key: _Key[_T] = ...
+    ) -> None: ...
+    def __new__(
+        cls, iterable: Optional[Iterable[_T]] = ..., key: _Key[_T] = ...
+    ) -> SortedKeyList[_T]: ...
+    @property
+    def key(self) -> Callable[[_T], Any]: ...
+    def clear(self) -> None: ...
+    def _clear(self) -> None: ...
+    def add(self, value: _T) -> None: ...
+    def _expand(self, pos: int) -> None: ...
+    def update(self, iterable: Iterable[_T]) -> None: ...
+    def _update(self, iterable: Iterable[_T]) -> None: ...
+    # NB: Must be T to be safely passed to self.func, yet base class imposes Any
+    def __contains__(self, value: _T) -> bool: ...  # type: ignore
+    def discard(self, value: _T) -> None: ...
+    def remove(self, value: _T) -> None: ...
+    def _delete(self, pos: int, idx: int) -> None: ...
+    def irange(
+        self,
+        minimum: Optional[_T] = ...,
+        maximum: Optional[_T] = ...,
+        inclusive: Tuple[bool, bool] = ...,
+        reverse: bool = ...,
+    ): ...
+    def irange_key(
+        self,
+        min_key: Optional[Any] = ...,
+        max_key: Optional[Any] = ...,
+        inclusive: Tuple[bool, bool] = ...,
+        reserve: bool = ...,
+    ): ...
+    def bisect_left(self, value: _T) -> int: ...
+    def bisect_right(self, value: _T) -> int: ...
+    def bisect(self, value: _T) -> int: ...
+    def bisect_key_left(self, key: Any) -> int: ...
+    def _bisect_key_left(self, key: Any) -> int: ...
+    def bisect_key_right(self, key: Any) -> int: ...
+    def _bisect_key_right(self, key: Any) -> int: ...
+    def bisect_key(self, key: Any) -> int: ...
+    def count(self, value: _T) -> int: ...
+    def copy(self: _SKL) -> _SKL: ...
+    def __copy__(self: _SKL) -> _SKL: ...
+    def index(
+        self, value: _T, start: Optional[int] = ..., stop: Optional[int] = ...
+    ) -> int: ...
+    def __add__(self: _SKL, other: Iterable[_T]) -> _SKL: ...
+    def __radd__(self: _SKL, other: Iterable[_T]) -> _SKL: ...
+    def __iadd__(self: _SKL, other: Iterable[_T]) -> _SKL: ...
+    def __mul__(self: _SKL, num: int) -> _SKL: ...
+    def __rmul__(self: _SKL, num: int) -> _SKL: ...
+    def __imul__(self: _SKL, num: int) -> _SKL: ...
+    def __repr__(self) -> str: ...
+    def _check(self) -> None: ...
+
+SortedListWithKey = SortedKeyList