don't add redundant types to unions to prevent unnecessary unknown/an…

…y errors
DetachHead · Jun 25, 2024 · ea424af · ea424af
1 parent 8ec869a
commit ea424af
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 11 deletions.
diff --git a/packages/pyright-internal/src/analyzer/codeFlowEngine.ts b/packages/pyright-internal/src/analyzer/codeFlowEngine.ts
@@ -906,7 +906,8 @@ export function getCodeFlowEngine(
                     }
                 }
 
-                const effectiveType = typesToCombine.length > 0 ? combineTypes(typesToCombine) : undefined;
+                const effectiveType =
+                    typesToCombine.length > 0 ? combineTypes(typesToCombine, undefined, evaluator) : undefined;
 
                 return setCacheEntry(branchNode, effectiveType, sawIncomplete);
             }

diff --git a/packages/pyright-internal/src/analyzer/operations.ts b/packages/pyright-internal/src/analyzer/operations.ts
@@ -656,7 +656,7 @@ export function getTypeOfBinaryOperation(
                 flags | EvaluatorFlags.ExpectingInstantiableType
             );
 
-            let newUnion = combineTypes([adjustedLeftType, adjustedRightType]);
+            let newUnion = combineTypes([adjustedLeftType, adjustedRightType], undefined, evaluator);
 
             const unionClass = evaluator.getUnionClassType();
             if (unionClass && isInstantiableClass(unionClass)) {

diff --git a/packages/pyright-internal/src/analyzer/types.ts b/packages/pyright-internal/src/analyzer/types.ts
@@ -12,6 +12,7 @@ import { Uri } from '../common/uri/uri';
 import { ArgumentNode, ExpressionNode, NameNode, ParameterCategory } from '../parser/parseNodes';
 import { ClassDeclaration, FunctionDeclaration, SpecialBuiltInClassDeclaration } from './declaration';
 import { Symbol, SymbolTable } from './symbol';
+import { TypeEvaluator } from './typeEvaluatorTypes';
 
 export const enum TypeCategory {
     // Name is not bound to a value of any type.
@@ -3237,7 +3238,6 @@ export function removeUnbound(type: Type): Type {
 
     return type;
 }
-
 export function removeFromUnion(type: Type, removeFilter: (type: Type) => boolean) {
     if (isUnion(type)) {
         const remainingTypes = type.subtypes.filter((t) => !removeFilter(t));
@@ -3269,7 +3269,7 @@ export function findSubtype(type: Type, filter: (type: UnionableType | NeverType
 // the same, only one is returned. If they differ, they
 // are combined into a UnionType. NeverTypes are filtered out.
 // If no types remain in the end, a NeverType is returned.
-export function combineTypes(subtypes: Type[], maxSubtypeCount?: number): Type {
+export function combineTypes(subtypes: Type[], maxSubtypeCount?: number, evaluator?: TypeEvaluator): Type {
     // Filter out any "Never" and "NoReturn" types.
     let sawNoReturn = false;
 
@@ -3352,21 +3352,33 @@ export function combineTypes(subtypes: Type[], maxSubtypeCount?: number): Type {
         return UnknownType.create();
     }
 
-    const newUnionType = UnionType.create();
+    let newUnionType = UnionType.create();
     if (typeAliasSources.size > 0) {
         newUnionType.typeAliasSources = typeAliasSources;
     }
 
     let hitMaxSubtypeCount = false;
 
     expandedTypes.forEach((subtype, index) => {
-        if (index === 0) {
-            UnionType.addType(newUnionType, subtype as UnionableType);
-        } else {
-            if (maxSubtypeCount === undefined || newUnionType.subtypes.length < maxSubtypeCount) {
-                _addTypeIfUnique(newUnionType, subtype as UnionableType);
+        if (!evaluator || !newUnionType.subtypes.length || !evaluator.assignType(newUnionType, subtype)) {
+            if (evaluator) {
+                const filteredType = removeFromUnion(newUnionType, (type) => evaluator.assignType(subtype, type));
+                if (isUnion(filteredType)) {
+                    newUnionType = filteredType;
+                } else {
+                    newUnionType = UnionType.create();
+                    UnionType.addType(newUnionType, filteredType as UnionableType);
+                }
+                newUnionType = isUnion(filteredType) ? filteredType : UnionType.create();
+            }
+            if (index === 0) {
+                UnionType.addType(newUnionType, subtype as UnionableType);
             } else {
-                hitMaxSubtypeCount = true;
+                if (maxSubtypeCount === undefined || newUnionType.subtypes.length < maxSubtypeCount) {
+                    _addTypeIfUnique(newUnionType, subtype as UnionableType);
+                } else {
+                    hitMaxSubtypeCount = true;
+                }
             }
         }
     });

diff --git a/packages/pyright-internal/src/tests/samples/typeNarrowingBased.py b/packages/pyright-internal/src/tests/samples/typeNarrowingBased.py
@@ -0,0 +1,16 @@
+from typing import Any, assert_type
+
+
+def foo(value: object):
+    print(value)
+    if isinstance(value, list):
+        _ = assert_type(value, list[Any])
+    _ = assert_type(value, object)
+
+def bar(value: object):
+    print(value)
+    if isinstance(value, list):
+        _ = assert_type(value, list[Any])
+    else:
+        _ = assert_type(value, object)
+    _ = assert_type(value, object)
diff --git a/packages/pyright-internal/src/tests/typeEvaluatorBased.test.ts b/packages/pyright-internal/src/tests/typeEvaluatorBased.test.ts
@@ -103,3 +103,12 @@ test('no reportUninitializedInstanceVariable on NamedTuple', () => {
         errors: [{ code: DiagnosticRule.reportUninitializedInstanceVariable, line: 6 }],
     });
 });
+
+test("useless type isn't added to union after if statement", () => {
+    const configOptions = new ConfigOptions(Uri.empty());
+    configOptions.diagnosticRuleSet.reportAssertTypeFailure = 'error';
+    const analysisResults = typeAnalyzeSampleFiles(['typeNarrowingBased.py'], configOptions);
+    validateResultsButBased(analysisResults, {
+        errors: [],
+    });
+});