From 75ca813a0428a8db2be7c356535e123fb5da4bac Mon Sep 17 00:00:00 2001
From: Matthew Hughes <matthewhughes934@gmail.com>
Date: Tue, 9 Aug 2022 20:13:15 +0100
Subject: [PATCH] Add type inference for `dict.keys` membership

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
7678f28f65cf6e878cacd5ba663e6e54347653ae).

GH: issue #13360
---
 mypy/checker.py                |  2 ++
 test-data/unit/pythoneval.test | 24 ++++++++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/mypy/checker.py b/mypy/checker.py
index 076f9e3763d9..b4d604dbfc49 100644
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -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).
diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test
index d7d20a923984..aedf7333354d 100644
--- a/test-data/unit/pythoneval.test
+++ b/test-data/unit/pythoneval.test
@@ -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]"