From f2328237d3ef7ad5b64af95cec5124458d79cfe3 Mon Sep 17 00:00:00 2001
From: Eric Traut <erictr@microsoft.com>
Date: Mon, 30 May 2022 00:55:22 -0700
Subject: [PATCH] Extended the "X in Y" type narrowing to support `dict`,
 `defaultdict` and `OrderedDict` classes.

---
 docs/type-concepts.md                                |  2 +-
 packages/pyright-internal/src/analyzer/typeGuards.ts |  8 ++++++--
 .../src/tests/samples/typeNarrowingIn1.py            | 12 ++++++++++++
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/docs/type-concepts.md b/docs/type-concepts.md
index f707dd608..aab49458b 100644
--- a/docs/type-concepts.md
+++ b/docs/type-concepts.md
@@ -178,7 +178,7 @@ In addition to assignment-based type narrowing, Pyright supports the following t
 * `x[I] == V` and `x[I] != V` (where I and V are literal expressions and x is a known-length tuple that is distinguished by the index indicated by I)
 * `x[I] is None` and `x[I] is not None` (where I is a literal expression and x is a known-length tuple that is distinguished by the index indicated by I)
 * `len(x) == L` and `len(x) != L` (where x is tuple and L is a literal integer)
-* `x in y` or `x not in y` (where y is instance of list, set, frozenset, deque, or tuple)
+* `x in y` or `x not in y` (where y is instance of list, set, frozenset, deque, tuple, dict, defaultdict, or OrderedDict)
 * `S in D` and `S not in D` (where S is a string literal and D is a TypedDict)
 * `isinstance(x, T)` (where T is a type or a tuple of types)
 * `issubclass(x, T)` (where T is a type or a tuple of types)
diff --git a/packages/pyright-internal/src/analyzer/typeGuards.ts b/packages/pyright-internal/src/analyzer/typeGuards.ts
index 390a40153..76df7d345 100644
--- a/packages/pyright-internal/src/analyzer/typeGuards.ts
+++ b/packages/pyright-internal/src/analyzer/typeGuards.ts
@@ -1336,11 +1336,15 @@ function narrowTypeForContains(evaluator: TypeEvaluator, referenceType: Type, co
 
     const builtInName = containerType.details.name;
 
-    if (!['list', 'set', 'frozenset', 'deque', 'tuple'].some((name) => name === builtInName)) {
+    if (
+        !['list', 'set', 'frozenset', 'deque', 'tuple', 'dict', 'defaultdict', 'OrderedDict'].some(
+            (name) => name === builtInName
+        )
+    ) {
         return referenceType;
     }
 
-    if (!containerType.typeArguments || containerType.typeArguments.length !== 1) {
+    if (!containerType.typeArguments || containerType.typeArguments.length < 1) {
         return referenceType;
     }
 
diff --git a/packages/pyright-internal/src/tests/samples/typeNarrowingIn1.py b/packages/pyright-internal/src/tests/samples/typeNarrowingIn1.py
index 875bfad36..1f85e2d16 100644
--- a/packages/pyright-internal/src/tests/samples/typeNarrowingIn1.py
+++ b/packages/pyright-internal/src/tests/samples/typeNarrowingIn1.py
@@ -83,3 +83,15 @@ def func4(val: str | None, container: list[str]):
         reveal_type(val, expected_text="str | None")
     else:
         reveal_type(val, expected_text="str")
+
+
+def func5(x: str | None, y: int | None, z: dict[str, str]):
+    if x in z:
+        reveal_type(x, expected_text="str")
+    else:
+        reveal_type(x, expected_text="str | None")
+
+    if y not in z:
+        reveal_type(y, expected_text="int | None")
+    else:
+        reveal_type(y, expected_text="Never")