wip

DetachHead · Jun 25, 2024 · a1c2a7f · a1c2a7f
1 parent ef1065e
commit a1c2a7f
Show file tree

Hide file tree

Showing 3 changed files with 20 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,9 @@ export function getCodeFlowEngine(
                     }
                 }
 
-                return setCacheEntry(branchNode, undefined, sawIncomplete);
+                const effectiveType = typesToCombine.length > 0 ? combineTypes(typesToCombine) : undefined;
+
+                return setCacheEntry(branchNode, effectiveType, sawIncomplete);
             }
 
             function getTypeFromLoopFlowNode(

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,8 @@ export function removeUnbound(type: Type): Type {
 
     return type;
 }
-
+export function removeFromUnion(type: UnionType, removeFilter: (type: Type) => boolean): UnionType;
+export function removeFromUnion(type: Type, removeFilter: (type: Type) => boolean): Type;
 export function removeFromUnion(type: Type, removeFilter: (type: Type) => boolean) {
     if (isUnion(type)) {
         const remainingTypes = type.subtypes.filter((t) => !removeFilter(t));
@@ -3269,7 +3271,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 +3354,26 @@ 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) {
+                newUnionType = removeFromUnion(newUnionType, (type) => evaluator.assignType(subtype, type));
+            }
+            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;
+                }
             }
         }
     });