Fix type narrowing for in operator with union containers and collection views

packages/pyright-internal/src/analyzer/typeGuards.ts

@@ -54,6 +54,7 @@ import {
     isParamSpec,
     isTypeSame,
     isTypeVar,
+    isUnion,
     isUnpackedTypeVarTuple,
     maxTypeRecursionCount,
     OverloadedType,
@@ -2199,23 +2200,72 @@ function narrowTypeForContainerType(
     });
 }
 
-export function getElementTypeForContainerNarrowing(containerType: Type) {
-    // We support contains narrowing only for certain built-in types that have been specialized.
-    const supportedContainers = ['list', 'set', 'frozenset', 'deque', 'tuple', 'dict', 'defaultdict', 'OrderedDict'];
-    if (!isClassInstance(containerType) || !ClassType.isBuiltIn(containerType, supportedContainers)) {
+export function getElementTypeForContainerNarrowing(containerType: Type): Type | undefined {
+    if (isUnion(containerType)) {
+        const elementTypes: Type[] = [];
+        for (const subtype of containerType.priv.subtypes) {
+            const elemType = getElementTypeForContainerNarrowing(subtype);
+            if (!elemType) {
+                return undefined;
+            }
+            elementTypes.push(elemType);
+        }
+        return combineTypes(elementTypes);
+    }
+
+    if (!isClassInstance(containerType)) {
+        return undefined;
+    }
+
+    // We support contains narrowing only for certain built-in and stdlib collection types that have been specialized.
+    const supportedContainers = [
+        'list',
+        'set',
+        'frozenset',
+        'deque',
+        'tuple',
+        'dict',
+        'defaultdict',
+        'OrderedDict',
+        'Sequence',
+        'MutableSequence',
+        'Set',
+        'AbstractSet',
+        'MutableSet',
+        'Collection',
+        'Container',
+        'Mapping',
+        'MutableMapping',
+        'KeysView',
+        'ValuesView',
+        'dict_keys',
+        'dict_values',
+    ];
+
+    const isSupported =
+        ClassType.isBuiltIn(containerType, supportedContainers) ||
+        ((containerType.shared.fullName.startsWith('_collections_abc.') ||
+            containerType.shared.fullName.startsWith('collections.abc.') ||
+            containerType.shared.fullName.startsWith('typing.')) &&
+            supportedContainers.includes(containerType.shared.name));
+
+    if (!isSupported) {
         return undefined;
     }
 
     if (!containerType.priv.typeArgs || containerType.priv.typeArgs.length < 1) {
         return undefined;
     }
 
-    let elementType = containerType.priv.typeArgs[0];
+    if (containerType.shared.name === 'dict_values' || containerType.shared.name === 'ValuesView') {
+        return containerType.priv.typeArgs.length > 1 ? containerType.priv.typeArgs[1] : containerType.priv.typeArgs[0];
+    }
+
     if (isTupleClass(containerType) && containerType.priv.tupleTypeArgs) {
-        elementType = combineTypes(containerType.priv.tupleTypeArgs.map((t) => t.type));
+        return combineTypes(containerType.priv.tupleTypeArgs.map((t) => t.type));
     }
 
-    return elementType;
+    return containerType.priv.typeArgs[0];
 }
 
 export function narrowTypeForContainerElementType(evaluator: TypeEvaluator, referenceType: Type, elementType: Type) {

packages/pyright-internal/src/tests/samples/typeNarrowingIn1.py

@@ -224,3 +224,32 @@ def func23[T: LiteralString](x: str, y: tuple[T, ...]) -> T:
     if x in y:
         return x
     raise ValueError(f"Invalid value {x!r}")
+
+
+def func24(val: str | None, container: list[str] | tuple[str, ...]):
+    if val in container:
+        reveal_type(val, expected_text="str")
+    else:
+        reveal_type(val, expected_text="str | None")
+
+
+def func25(val: str | int | None, container: list[str] | set[int]):
+    if val in container:
+        reveal_type(val, expected_text="str | int")
+    else:
+        reveal_type(val, expected_text="str | int | None")
+
+
+def func26(k: str | int, d: dict[str, float]):
+    if k in d.keys():
+        reveal_type(k, expected_text="str")
+    else:
+        reveal_type(k, expected_text="str | int")
+
+
+def func27(v: float | str, d: dict[str, float]):
+    if v in d.values():
+        reveal_type(v, expected_text="float")
+    else:
+        reveal_type(v, expected_text="float | str")
+