diff --git a/packages/pyright-internal/src/analyzer/patternMatching.ts b/packages/pyright-internal/src/analyzer/patternMatching.ts index e411429168e3..c212532489e7 100644 --- a/packages/pyright-internal/src/analyzer/patternMatching.ts +++ b/packages/pyright-internal/src/analyzer/patternMatching.ts @@ -469,58 +469,90 @@ function narrowTypeBasedOnMappingPattern( return combineTypes(mappingInfo.filter((m) => !m.isDefinitelyMapping).map((m) => m.subtype)); } - if (pattern.d.entries.length !== 1 || pattern.d.entries[0].nodeType !== ParseNodeType.PatternMappingKeyEntry) { - return type; - } - // Handle the case where the type is a union that includes a TypedDict with - // a field discriminated by a literal. - const keyPattern = pattern.d.entries[0].d.keyPattern; - const valuePattern = pattern.d.entries[0].d.valuePattern; - if ( - keyPattern.nodeType !== ParseNodeType.PatternLiteral || - valuePattern.nodeType !== ParseNodeType.PatternAs || - !valuePattern.d.orPatterns.every((orPattern) => orPattern.nodeType === ParseNodeType.PatternLiteral) - ) { - return type; - } + // fields discriminated by literals or match patterns. + const keyEntryInfos: { keyValue: string; valueTypes?: Type[] }[] = []; + let hasUnsupportedKeyPattern = false; - const keyType = evaluator.getTypeOfExpression(keyPattern.d.expr).type; + for (const entry of pattern.d.entries) { + if (entry.nodeType === ParseNodeType.PatternMappingExpandEntry) { + continue; + } + if (entry.nodeType !== ParseNodeType.PatternMappingKeyEntry) { + hasUnsupportedKeyPattern = true; + break; + } - // The key type must be a str literal. - if ( - !isClassInstance(keyType) || - !ClassType.isBuiltIn(keyType, 'str') || - keyType.priv.literalValue === undefined - ) { - return type; + const keyPattern = entry.d.keyPattern; + if (keyPattern.nodeType !== ParseNodeType.PatternLiteral) { + hasUnsupportedKeyPattern = true; + break; + } + + const keyType = evaluator.getTypeOfExpression(keyPattern.d.expr).type; + if ( + !isClassInstance(keyType) || + !ClassType.isBuiltIn(keyType, 'str') || + keyType.priv.literalValue === undefined + ) { + hasUnsupportedKeyPattern = true; + break; + } + const keyValue = keyType.priv.literalValue as string; + + const valuePattern = entry.d.valuePattern; + let valueTypes: Type[] | undefined; + + if ( + valuePattern.nodeType === ParseNodeType.PatternAs && + valuePattern.d.orPatterns.every((orPattern) => orPattern.nodeType === ParseNodeType.PatternLiteral) + ) { + valueTypes = valuePattern.d.orPatterns.map( + (orPattern) => evaluator.getTypeOfExpression((orPattern as PatternLiteralNode).d.expr).type + ); + } else if ( + valuePattern.nodeType === ParseNodeType.PatternAs && + valuePattern.d.orPatterns.length === 1 && + valuePattern.d.orPatterns[0].nodeType === ParseNodeType.PatternCapture + ) { + valueTypes = undefined; + } else { + hasUnsupportedKeyPattern = true; + break; + } + + keyEntryInfos.push({ keyValue, valueTypes }); } - const keyValue = keyType.priv.literalValue as string; - const valueTypes = valuePattern.d.orPatterns.map( - (orPattern) => evaluator.getTypeOfExpression((orPattern as PatternLiteralNode).d.expr).type - ); + if (hasUnsupportedKeyPattern || keyEntryInfos.length === 0) { + return type; + } return mapSubtypes(type, (subtype) => { if (isClassInstance(subtype) && ClassType.isTypedDictClass(subtype)) { const typedDictMembers = getTypedDictMembersForClass(evaluator, subtype, /* allowNarrowed */ true); - const member = typedDictMembers.knownItems.get(keyValue); - if (member && (member.isRequired || member.isProvided) && isClassInstance(member.valueType)) { - const memberValueType = member.valueType; + const matchesAllEntries = keyEntryInfos.every((entryInfo) => { + const member = typedDictMembers.knownItems.get(entryInfo.keyValue); + if (!member || (!member.isRequired && !member.isProvided) || !isClassInstance(member.valueType)) { + return false; + } - // If there's at least one literal value pattern that matches - // the literal type of the member, we can eliminate this type. - if ( - valueTypes.some( - (valueType) => - isClassInstance(valueType) && - ClassType.isSameGenericClass(valueType, memberValueType) && - valueType.priv.literalValue === memberValueType.priv.literalValue - ) - ) { - return undefined; + if (entryInfo.valueTypes === undefined) { + return true; } + + const memberValueType = member.valueType; + return entryInfo.valueTypes.some( + (valueType) => + isClassInstance(valueType) && + ClassType.isSameGenericClass(valueType, memberValueType) && + valueType.priv.literalValue === memberValueType.priv.literalValue + ); + }); + + if (matchesAllEntries) { + return undefined; } } diff --git a/packages/pyright-internal/src/tests/samples/matchMapping1.py b/packages/pyright-internal/src/tests/samples/matchMapping1.py index 9359dfa0f87e..fbcbcde2e6fc 100644 --- a/packages/pyright-internal/src/tests/samples/matchMapping1.py +++ b/packages/pyright-internal/src/tests/samples/matchMapping1.py @@ -1,6 +1,7 @@ # This sample tests type checking for match statements (as # described in PEP 634) that contain mapping patterns. +from enum import Enum from typing import Literal, TypedDict from typing_extensions import NotRequired # pyright: ignore[reportMissingModuleSource] @@ -150,5 +151,56 @@ def test_not_required_narrowing(subj: TD1) -> None: # This should generate an error. print(subj["v1"]) + print(subj["v2"]) print(subj["v3"]) + + +class MsgA(TypedDict): + v: Literal[1] + kind: Literal["a"] + data_a: int + + +class MsgB(TypedDict): + v: Literal[1] + kind: Literal["b"] + data_b: str + + +def test_negative_narrowing3(msg: MsgA | MsgB) -> None: + match msg: + case {"v": 1, "kind": "a"}: + reveal_type(msg, expected_text="MsgA") + case _: + reveal_type(msg, expected_text="MsgB") + + +class Color(Enum): + RED = 1 + BLUE = 2 + + +class RedMsg(TypedDict): + color: Literal[Color.RED] + + +class BlueMsg(TypedDict): + color: Literal[Color.BLUE] + + +def test_value_pattern_negative_narrowing(msg: RedMsg | BlueMsg) -> None: + match msg: + case {"color": Color.RED}: + reveal_type(msg, expected_text="RedMsg") + case _: + reveal_type(msg, expected_text="RedMsg | BlueMsg") + + +def test_capture_pattern_negative_narrowing(msg: MsgA | MsgB) -> None: + match msg: + case {"v": 1, "kind": x}: + reveal_type(msg, expected_text="MsgA | MsgB") + reveal_type(x, expected_text="Literal['a', 'b']") + case _: + reveal_type(msg, expected_text="Never")