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
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.