Short-circuit isTypeVarSame subtype checks

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

@@ -336,25 +336,20 @@ export function isTypeVarSame(type1: TypeVarType, type2: Type) {
         return false;
     }
 
-    let isCompatible = true;
-    doForEachSubtype(type2, (subtype) => {
-        if (!isCompatible) {
-            return;
-        }
-
+    return allSubtypes(type2, (subtype) => {
         if (!isTypeSame(type1, subtype)) {
             const conditions = getTypeCondition(subtype);
 
             if (
                 !conditions ||
                 !conditions.some((condition) => condition.typeVar.priv.nameWithScope === type1.priv.nameWithScope)
             ) {
-                isCompatible = false;
+                return false;
             }
         }
-    });
 
-    return isCompatible;
+        return true;
+    });
 }
 
 export function makeInferenceContext(
@@ -796,12 +791,10 @@ export function someSubtypes(type: Type, callback: (type: Type) => boolean): boo
 
 export function allSubtypes(type: Type, callback: (type: Type) => boolean): boolean {
     if (isUnion(type)) {
-        return type.priv.subtypes.every((subtype) => {
-            callback(subtype);
-        });
-    } else {
-        return callback(type);
+        return type.priv.subtypes.every(callback);
     }
+
+    return callback(type);
 }
 
 export function doForEachSignature(

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

@@ -8,7 +8,7 @@
 
 import * as assert from 'assert';
 
-import { transformTypePair } from '../analyzer/typeUtils';
+import { allSubtypes, someSubtypes, transformTypePair } from '../analyzer/typeUtils';
 import {
     AnyType,
     ClassType,
@@ -22,6 +22,7 @@ import {
     Type,
     TypeVarScopeType,
     TypeVarType,
+    UnionableType,
     UnionType,
     UnknownType,
     Variance,
@@ -164,3 +165,84 @@ function replacePair(source: Type, target: Type) {
     return (sourceNode: Type, targetNode: Type) =>
         sourceNode === source && targetNode === target ? targetNode : undefined;
 }
+
+test('AllSubtypes', () => {
+    const unionType = createUnion(createClassType('A'), createClassType('B'), createClassType('C'));
+    const visitedSubtypes: Type[] = [];
+
+    const result = allSubtypes(unionType, (subtype) => {
+        visitedSubtypes.push(subtype);
+        return visitedSubtypes.length < 2;
+    });
+
+    assert.strictEqual(result, false);
+    assert.strictEqual(visitedSubtypes.length, 2);
+
+    assert.strictEqual(
+        allSubtypes(unionType, () => {
+            return true;
+        }),
+        true
+    );
+
+    const singleType = createClassType('D');
+    assert.strictEqual(
+        allSubtypes(singleType, (subtype) => {
+            assert.strictEqual(subtype, singleType);
+            return true;
+        }),
+        true
+    );
+});
+
+test('SomeSubtypes', () => {
+    const unionType = createUnion(createClassType('A'), createClassType('B'), createClassType('C'));
+    const visitedSubtypes: Type[] = [];
+
+    const result = someSubtypes(unionType, (subtype) => {
+        visitedSubtypes.push(subtype);
+        return visitedSubtypes.length === 2;
+    });
+
+    assert.strictEqual(result, true);
+    assert.strictEqual(visitedSubtypes.length, 2);
+
+    assert.strictEqual(
+        someSubtypes(unionType, () => {
+            return false;
+        }),
+        false
+    );
+
+    const singleType = createClassType('D');
+    assert.strictEqual(
+        someSubtypes(singleType, (subtype) => {
+            assert.strictEqual(subtype, singleType);
+            return false;
+        }),
+        false
+    );
+});
+
+function createClassType(name: string, flags = ClassTypeFlags.None) {
+    const classType = ClassType.createInstantiable(
+        name,
+        name,
+        '',
+        Uri.empty(),
+        flags,
+        0,
+        /* declaredMetaclass*/ undefined,
+        /* effectiveMetaclass */ undefined
+    );
+    classType.shared.mro.push(classType);
+    return classType;
+}
+
+function createUnion(...subtypes: UnionableType[]) {
+    const unionType = UnionType.create();
+    subtypes.forEach((subtype) => {
+        UnionType.addType(unionType, subtype);
+    });
+    return unionType;
+}