Add support for type narrowing of NamedTuple fields with "is None" (#3271)

* Add support for type narrowing of NamedTuple fields with is None.

* Add description for `x.y is None` narrowing

Fix capitalization for filename of sample in test.

docs/type-concepts.md

@@ -171,6 +171,7 @@ In addition to assignment-based type narrowing, Pyright supports the following t
 * `type(x) is T` and `type(x) is not T`
 * `x is E` and `x is not E` (where E is a literal enum or bool)
 * `x == L` and `x != L` (where L is a literal expression)
+* `x.y is None` and `x.y is not None` (where x is a type that is distinguished by a field with a None)
 * `x.y is E` and `x.y is not E` (where E is a literal enum or bool and x is a type that is distinguished by a field with a literal type)
 * `x.y == L` and `x.y != L` (where L is a literal expression and x is a type that is distinguished by a field with a literal type)
 * `x[K] == V` and `x[K] != V` (where K and V are literal expressions and x is a type that is distinguished by a TypedDict field with a literal type)

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

@@ -65,6 +65,7 @@ import {
     convertToInstance,
     convertToInstantiable,
     doForEachSubtype,
+    getSpecializedTupleType,
     getTypeCondition,
     getTypeVarScopeId,
     isLiteralType,
@@ -355,6 +356,25 @@ export function getTypeNarrowingCallback(
                     };
                 }
             }
+
+            // Look for X.Y is None or X.Y is not None
+            // These are commonly-used patterns used in control flow.
+            if (
+                testExpression.leftExpression.nodeType === ParseNodeType.MemberAccess &&
+                ParseTreeUtils.isMatchingExpression(reference, testExpression.leftExpression.leftExpression) &&
+                testExpression.rightExpression.nodeType === ParseNodeType.Constant &&
+                testExpression.rightExpression.constType === KeywordType.None
+            ) {
+                const memberName = testExpression.leftExpression.memberName;
+                return (type: Type) => {
+                    return narrowTypeForDiscriminatedFieldNoneComparison(
+                        evaluator,
+                        type,
+                        memberName.value,
+                        adjIsPositiveTest
+                    );
+                };
+            }
         }
 
         if (testExpression.operator === OperatorType.In) {
@@ -707,24 +727,20 @@ function narrowTypeForTruthiness(evaluator: TypeEvaluator, type: Type, isPositiv
 }
 
 // Handle type narrowing for expressions of the form "a[I] is None" and "a[I] is not None" where
-// I is an integer and a is a union of Tuples with known lengths and entry types.
+// I is an integer and a is a union of Tuples and NamedTuples with known lengths and entry types.
 function narrowTupleTypeForIsNone(evaluator: TypeEvaluator, type: Type, isPositiveTest: boolean, indexValue: number) {
     return evaluator.mapSubtypesExpandTypeVars(type, /* conditionFilter */ undefined, (subtype) => {
-        if (
-            !isClassInstance(subtype) ||
-            !isTupleClass(subtype) ||
-            isUnboundedTupleClass(subtype) ||
-            !subtype.tupleTypeArguments
-        ) {
+        const tupleType = getSpecializedTupleType(subtype);
+        if (!tupleType || isUnboundedTupleClass(tupleType) || !tupleType.tupleTypeArguments) {
             return subtype;
         }
 
-        const tupleLength = subtype.tupleTypeArguments.length;
+        const tupleLength = tupleType.tupleTypeArguments.length;
         if (indexValue < 0 || indexValue >= tupleLength) {
             return subtype;
         }
 
-        const typeOfEntry = evaluator.makeTopLevelTypeVarsConcrete(subtype.tupleTypeArguments[indexValue].type);
+        const typeOfEntry = evaluator.makeTopLevelTypeVarsConcrete(tupleType.tupleTypeArguments[indexValue].type);
 
         if (isPositiveTest) {
             if (!evaluator.canAssignType(typeOfEntry, NoneType.createInstance())) {
@@ -1459,6 +1475,41 @@ function narrowTypeForDiscriminatedFieldComparison(
     return narrowedType;
 }
 
+// Attempts to narrow a type based on a comparison (equal or not equal)
+// between a discriminating field that has a declared None type to a
+// None.
+function narrowTypeForDiscriminatedFieldNoneComparison(
+    evaluator: TypeEvaluator,
+    referenceType: Type,
+    memberName: string,
+    isPositiveTest: boolean
+): Type {
+    return mapSubtypes(referenceType, (subtype) => {
+        let memberInfo: ClassMember | undefined;
+        if (isClassInstance(subtype)) {
+            memberInfo = lookUpObjectMember(subtype, memberName);
+        } else if (isInstantiableClass(subtype)) {
+            memberInfo = lookUpClassMember(subtype, memberName);
+        }
+
+        if (memberInfo && memberInfo.isTypeDeclared) {
+            const memberType = evaluator.getTypeOfMember(memberInfo);
+
+            if (isPositiveTest) {
+                if (!evaluator.canAssignType(memberType, NoneType.createInstance())) {
+                    return undefined;
+                }
+            } else {
+                if (isNoneInstance(memberType)) {
+                    return undefined;
+                }
+            }
+        }
+
+        return subtype;
+    });
+}
+
 // Attempts to narrow a type based on a "type(x) is y" or "type(x) is not y" check.
 function narrowTypeForTypeIs(type: Type, classType: ClassType, isPositiveTest: boolean) {
     return mapSubtypes(type, (subtype) => {

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

@@ -0,0 +1,35 @@
+# This sample tests the type narrowing case for unions of NamedTuples
+# where one or more of the entries is tested against type None by attribute.
+
+from typing import NamedTuple, Union
+
+IntFirst = NamedTuple("IntFirst", [
+    ("first", int),
+    ("second", None),
+])
+
+StrSecond = NamedTuple("StrSecond", [
+    ("first", None),
+    ("second", str),
+])
+
+def func1(a: Union[IntFirst, StrSecond]) -> IntFirst:
+    if a.second is None:
+        reveal_type(a, expected_text="IntFirst")
+        return a
+    else:
+        reveal_type(a, expected_text="StrSecond")
+        raise ValueError()
+
+
+UnionFirst = NamedTuple("UnionFirst", [
+    ("first", Union[None, int]),
+    ("second", None),
+])
+
+def func2(a: Union[UnionFirst, StrSecond]):
+    if a.first is None:
+        reveal_type(a, expected_text="UnionFirst | StrSecond")
+    else:
+        reveal_type(a, expected_text="UnionFirst")
+

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

@@ -0,0 +1,35 @@
+# This sample tests the type narrowing case for unions of NamedTuples
+# where one or more of the entries is tested against type None by index.
+
+from typing import NamedTuple, Union
+
+IntFirst = NamedTuple("IntFirst", [
+    ("first", int),
+    ("second", None),
+])
+
+StrSecond = NamedTuple("StrSecond", [
+    ("first", None),
+    ("second", str),
+])
+
+def func1(a: Union[IntFirst, StrSecond]) -> IntFirst:
+    if a[1] is None:
+        reveal_type(a, expected_text="IntFirst")
+        return a
+    else:
+        reveal_type(a, expected_text="StrSecond")
+        raise ValueError()
+
+
+UnionFirst = NamedTuple("UnionFirst", [
+    ("first", Union[None, int]),
+    ("second", None),
+])
+
+def func2(a: Union[UnionFirst, StrSecond]):
+    if a[0] is None:
+        reveal_type(a, expected_text="UnionFirst | StrSecond")
+    else:
+        reveal_type(a, expected_text="UnionFirst")
+

packages/pyright-internal/src/tests/typeEvaluator1.test.ts

@@ -302,6 +302,18 @@ test('TypeNarrowingIsNoneTuple1', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('TypeNarrowingIsNoneNamedTuple1', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingIsNoneNamedTuple1.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
+test('TypeNarrowingIsNoneNamedTuple2', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingIsNoneNamedTuple2.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('TypeNarrowingLiteral1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingLiteral1.py']);