diff --git a/packages/pyright-internal/src/analyzer/typeGuards.ts b/packages/pyright-internal/src/analyzer/typeGuards.ts index 4d6605ea2162..059926cf0796 100644 --- a/packages/pyright-internal/src/analyzer/typeGuards.ts +++ b/packages/pyright-internal/src/analyzer/typeGuards.ts @@ -54,6 +54,7 @@ import { isParamSpec, isTypeSame, isTypeVar, + isUnion, isUnpackedTypeVarTuple, maxTypeRecursionCount, OverloadedType, @@ -2199,10 +2200,56 @@ 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; } @@ -2210,12 +2257,15 @@ export function getElementTypeForContainerNarrowing(containerType: Type) { 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) { diff --git a/packages/pyright-internal/src/tests/samples/typeNarrowingIn1.py b/packages/pyright-internal/src/tests/samples/typeNarrowingIn1.py index de42085d901c..13173283ee96 100644 --- a/packages/pyright-internal/src/tests/samples/typeNarrowingIn1.py +++ b/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") +