python type of SortedDict can not be defined

Question

How can I define the type of SortedDict?

from typing import Dict
from sortedcontainers import SortedDict

x: Dict[int, str] = {2: "two"}
y: SortedDict[int, str] = {3: "three"} # <--- this doesn't work ...

I realize that the exact name SortedDict will not do as this is the actual class, but what will work?

$ python3.8 example.py 
Traceback (most recent call last):
  File "example.py", line 5, in <module>
    y: SortedDict[int, str] = {3: "three"}
TypeError: 'type' object is not subscriptable

Answer 1

Since the original packages has no type information, what you can do is make your own interface stub like this:

from typing import Any, Dict, Hashable, Iterable, Optional, Protocol, Tuple, TypeVar

class Comparable(Protocol):
    def __lt__(self, other: Any) -> bool: ...

K = TypeVar("K", bound=Hashable)
V = TypeVar("V", bound=Comparable)

class SortedDict(Dict[K, V]):
    def bisect_left(self, value: V) -> int: ...
    def bisect_right(self, value: V) -> int: ...
    def index(
        self, value: V, start: Optional[int] = None, stop: Optional[int] = None
    ) -> int: ...
    def irange(
        self,
        minimum: Optional[V] = None,
        maximum: Optional[V] = None,
        inclusive: Tuple[bool, bool] = (True, True),
        reverse: bool = False,
    ) -> Iterable[V]: ...
    def islice(
        self,
        start: Optional[int] = None,
        stop: Optional[int] = None,
        reverse: bool = False,
    ) -> Iterable[V]: ...

(Notice I made the value type "sortable" so the type checker verifies you don't pass it a list).

Save this as sortedcontainers.pyi and then set mypy's search path to it:

MYPYPATH=/path/to/interfaces mypy <args...>

Now you can do the following:

# example.py
from __future__ import annotations

from typing import Dict, Optional

from sortedcontainers import SortedDict

x: Dict[int, str] = {2: "two"}
y = SortedDict({3: "three"})
def foo(z: SortedDict[int, str]) -> Optional[str]:
    if 3 in z:
        z_3 = z[3]
        if isinstance(z_3, str):
            return z_3
    return None
t = foo(x) # <--- x is not a SortedDict
if t:
    print(t)

If you check it, you'll see that mypy catches the error:

$ mypy example.py
example.py:16: error: Argument 1 to "foo" has incompatible type "Dict[int, str]"; expected "SortedDict[int, str]"
Found 1 error in 1 file (checked 1 source file)

Hope this helped.