Skip to content

Commit

Permalink
Add type inference for dict.keys membership
Browse files Browse the repository at this point in the history
Also for containership checks on `typing.KeysView.` This is to bring the
following cases into alignment:

    from typing import KeysView

    key: str | None
    d: dict[str, int]
    kv: KeysView[str]

    if key in d:
        # type of 'key' is inferred to be 'str'
        ...

    if key in d.keys():
        # type of 'key' should be inferred to be 'str'
        ...

    if key in kv:
        # type of 'key' should be inferred to be 'str'
        ...

Before this change only the first `if` would narrow the type of `key`.

I've just added a test under `test-data/unit/pythoneval.test` as
`test-data/unit/fixtures/dict.pyi` doesn't include `dict.keys`, and
adding it requires importing `dict_keys` from `_collections_abc` in
those stubs, which then requires adding `_collections_abc.pyi` stubs,
which would have to be complete since e.g.
`testGenericAliasCollectionsABCReveal` expects most of the types in
those stubs to be defined (this is the same approach as
7678f28).

GH: issue #13360
  • Loading branch information
matthewhughes934 authored and Matt Hughes committed Aug 23, 2022
1 parent 4bb7688 commit 75ca813
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 0 deletions.
2 changes: 2 additions & 0 deletions mypy/checker.py
Original file line number Diff line number Diff line change
Expand Up @@ -6329,6 +6329,8 @@ def builtin_item_type(tp: Type) -> Type | None:
"builtins.dict",
"builtins.set",
"builtins.frozenset",
"_collections_abc.dict_keys",
"typing.KeysView",
]:
if not tp.args:
# TODO: fix tuple in lib-stub/builtins.pyi (it should be generic).
Expand Down
24 changes: 24 additions & 0 deletions test-data/unit/pythoneval.test
Original file line number Diff line number Diff line change
Expand Up @@ -1636,3 +1636,27 @@ foo("")
foo(list(""))
foo(list((list(""), "")))
[out]

[case testNarrowTypeForDictKeys]
# flags: --strict-optional
from typing import Dict, KeysView, Optional

d: Dict[str, int]
key: Optional[str]
if key in d.keys():
reveal_type(key)
else:
reveal_type(key)

kv: KeysView[str]
k: Optional[str]
if k in kv:
reveal_type(k)
else:
reveal_type(k)

[out]
_testNarrowTypeForDictKeys.py:7: note: Revealed type is "builtins.str"
_testNarrowTypeForDictKeys.py:9: note: Revealed type is "Union[builtins.str, None]"
_testNarrowTypeForDictKeys.py:14: note: Revealed type is "builtins.str"
_testNarrowTypeForDictKeys.py:16: note: Revealed type is "Union[builtins.str, None]"

0 comments on commit 75ca813

Please sign in to comment.