diff --git a/pyrsistent/__init__.pyi b/pyrsistent/__init__.pyi index 5909f79..af2c69c 100644 --- a/pyrsistent/__init__.pyi +++ b/pyrsistent/__init__.pyi @@ -43,6 +43,10 @@ T = TypeVar('T') KT = TypeVar('KT') VT = TypeVar('VT') +@overload +def pmap(initial: Mapping[KT, VT] = {}, pre_size: int = 0) -> PMap[KT, VT]: ... +@overload +def pmap(initial: Iterable[Tuple[KT, VT]] = {}, pre_size: int = 0) -> PMap[KT, VT]: ... def pmap(initial: Union[Mapping[KT, VT], Iterable[Tuple[KT, VT]]] = {}, pre_size: int = 0) -> PMap[KT, VT]: ... def m(**kwargs: VT) -> PMap[str, VT]: ... diff --git a/tests/typing_test.py b/tests/typing_test.py new file mode 100644 index 0000000..471cbe7 --- /dev/null +++ b/tests/typing_test.py @@ -0,0 +1,6 @@ +# Check that the inferred types are as expected. + +from typing_extensions import assert_type +from pyrsistent import PMap, pmap + +assert_type(pmap(dict[tuple[int, str], float]()), PMap[tuple[int, str], float]) diff --git a/tox.ini b/tox.ini index 834a2ba..612999a 100644 --- a/tox.ini +++ b/tox.ini @@ -25,6 +25,7 @@ commands = pytest deps = pytest hypothesis + typing_extensions # Specifying the following tests like this is very non-DRY but I have no better solution right now. [testenv:coverage-py38]