Further simplified TypeVar constraint tracking and solving. Removed precomputed "lowerBoundNoLiterals" from tracker. (#8656)

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

@@ -45,6 +45,7 @@ import {
 } from './types';
 import {
     addConditionToType,
+    applySolvedTypeVars,
     AssignTypeFlags,
     buildSolutionFromSpecializedClass,
     convertToInstance,
@@ -212,6 +213,24 @@ export function solveConstraints(
     return new ConstraintSolution(solutionSets);
 }
 
+// Applies solved TypeVars from one context to this context.
+export function applySourceSolutionToConstraints(constraints: ConstraintTracker, srcSolution: ConstraintSolution) {
+    if (srcSolution.isEmpty()) {
+        return;
+    }
+
+    constraints.doForEachConstraintSet((constraintSet) => {
+        constraintSet.getTypeVars().forEach((entry) => {
+            constraintSet.setBounds(
+                entry.typeVar,
+                entry.lowerBound ? applySolvedTypeVars(entry.lowerBound, srcSolution) : undefined,
+                entry.upperBound ? applySolvedTypeVars(entry.upperBound, srcSolution) : undefined,
+                entry.retainLiterals
+            );
+        });
+    });
+}
+
 export function solveConstraintSet(
     evaluator: TypeEvaluator,
     constraintSet: ConstraintSet,
@@ -229,38 +248,6 @@ export function solveConstraintSet(
     return solutionSet;
 }
 
-// Updates the lower and upper bounds for a type variable. It also calculates the
-// lowerBoundNoLiterals, which is a variant of the lower bound that has
-// literals stripped. By default, the constraint solver always uses the "no literals"
-// type in its solutions unless the version with literals is required to satisfy
-// the upper bound.
-export function updateTypeVarType(
-    evaluator: TypeEvaluator,
-    constraints: ConstraintTracker,
-    destType: TypeVarType,
-    lowerBound: Type | undefined,
-    upperBound: Type | undefined,
-    forceRetainLiterals = false
-) {
-    let lowerBoundNoLiterals: Type | undefined;
-
-    if (lowerBound && !forceRetainLiterals) {
-        const strippedLiteral = isTypeVarTuple(destType)
-            ? stripLiteralValueForUnpackedTuple(evaluator, lowerBound)
-            : evaluator.stripLiteralValue(lowerBound);
-
-        // Strip the literals from the lower bound and see if it is still
-        // narrower than the upper bound.
-        if (strippedLiteral !== lowerBound) {
-            if (!upperBound || evaluator.assignType(upperBound, strippedLiteral)) {
-                lowerBoundNoLiterals = strippedLiteral;
-            }
-        }
-    }
-
-    constraints.setBounds(destType, lowerBound, lowerBoundNoLiterals, upperBound);
-}
-
 // In cases where the expected type is a specialized base class of the
 // source type, we need to determine which type arguments in the derived
 // class will make it compatible with the specialized base class. This method
@@ -277,7 +264,7 @@ export function addConstraintsForExpectedType(
 ): boolean {
     if (isAny(expectedType)) {
         type.shared.typeParams.forEach((typeParam) => {
-            updateTypeVarType(evaluator, constraints, typeParam, expectedType, expectedType);
+            constraints.setBounds(typeParam, expectedType, expectedType);
         });
         return true;
     }
@@ -320,9 +307,7 @@ export function addConstraintsForExpectedType(
             if (typeArgValue) {
                 const variance = TypeVarType.getVariance(typeParam);
 
-                updateTypeVarType(
-                    evaluator,
-                    constraints,
+                constraints.setBounds(
                     typeParam,
                     variance === Variance.Covariant ? undefined : typeArgValue,
                     variance === Variance.Contravariant ? undefined : typeArgValue
@@ -438,9 +423,7 @@ export function addConstraintsForExpectedType(
                             typeArgValue = UnknownType.create();
                         }
 
-                        updateTypeVarType(
-                            evaluator,
-                            constraints,
+                        constraints.setBounds(
                             targetTypeVar,
                             variance === Variance.Covariant ? undefined : typeArgValue,
                             variance === Variance.Contravariant ? undefined : typeArgValue
@@ -472,19 +455,22 @@ export function applyUnificationVars(evaluator: TypeEvaluator, constraints: Cons
                 const newLowerBound = entry.lowerBound
                     ? applyUnificationVarsToType(evaluator, entry.lowerBound, constraintSet)
                     : undefined;
-                const newLowerBoundNoLiterals = entry.lowerBoundNoLiterals
-                    ? applyUnificationVarsToType(evaluator, entry.lowerBoundNoLiterals, constraintSet)
-                    : undefined;
                 const newUpperBound = entry.upperBound
                     ? applyUnificationVarsToType(evaluator, entry.upperBound, constraintSet)
                     : undefined;
 
-                constraintSet.setBounds(entry.typeVar, newLowerBound, newLowerBoundNoLiterals, newUpperBound);
+                constraintSet.setBounds(entry.typeVar, newLowerBound, newUpperBound, entry.retainLiterals);
             }
         });
     });
 }
 
+function stripLiteralsForLowerBound(evaluator: TypeEvaluator, typeVar: TypeVarType, lowerBound: Type) {
+    return isTypeVarTuple(typeVar)
+        ? stripLiteralValueForUnpackedTuple(evaluator, lowerBound)
+        : evaluator.stripLiteralValue(lowerBound);
+}
+
 function getTypeVarType(
     evaluator: TypeEvaluator,
     constraintSet: ConstraintSet,
@@ -514,10 +500,35 @@ function getTypeVarType(
         return entry.lowerBound;
     }
 
-    // Prefer the lower bound with no literals. It will be undefined
-    // if the literal type couldn't be widened due to constraints imposed
-    // by the upper bound.
-    return entry.lowerBoundNoLiterals ?? entry.lowerBound ?? entry.upperBound;
+    let result: Type | undefined;
+
+    let lowerBound = entry.lowerBound;
+    if (lowerBound) {
+        if (!entry.retainLiterals) {
+            const lowerNoLiterals = stripLiteralsForLowerBound(evaluator, typeVar, lowerBound);
+
+            // If we can widen the lower bound to a non-literal type without
+            // exceeding the upper bound, use the widened type.
+            if (lowerNoLiterals !== lowerBound) {
+                if (!entry.upperBound || evaluator.assignType(entry.upperBound, lowerNoLiterals)) {
+                    if (TypeVarType.hasConstraints(typeVar)) {
+                        // Does it still match a value constraint?
+                        if (typeVar.shared.constraints.some((constraint) => isTypeSame(lowerNoLiterals, constraint))) {
+                            lowerBound = lowerNoLiterals;
+                        }
+                    } else {
+                        lowerBound = lowerNoLiterals;
+                    }
+                }
+            }
+        }
+
+        result = lowerBound;
+    } else {
+        result = entry.upperBound;
+    }
+
+    return result;
 }
 
 // Handles an assignment to a TypeVar that is "bound" rather than "free".
@@ -723,8 +734,8 @@ function assignUnconstrainedTypeVar(
             // If this is an invariant context and there is currently no upper bound
             // established, use the "no literals" version of the lower bound rather
             // than a version that has literals.
-            if (!newUpperTypeBound && isInvariant && curEntry?.lowerBoundNoLiterals) {
-                newLowerBound = curEntry.lowerBoundNoLiterals;
+            if (!newUpperTypeBound && isInvariant && curEntry && !curEntry.retainLiterals) {
+                newLowerBound = stripLiteralsForLowerBound(evaluator, destType, curLowerBound);
             }
         } else {
             if (
@@ -820,8 +831,8 @@ function assignUnconstrainedTypeVar(
                     // If this is an invariant context and there is currently no upper bound
                     // established, use the "no literals" version of the lower bound rather
                     // than a version that has literals.
-                    if (!newUpperTypeBound && isInvariant && curEntry?.lowerBoundNoLiterals) {
-                        curLowerBound = curEntry.lowerBoundNoLiterals;
+                    if (!newUpperTypeBound && isInvariant && curEntry && !curEntry.retainLiterals) {
+                        curLowerBound = stripLiteralsForLowerBound(evaluator, destType, curLowerBound);
                     }
 
                     let curSolvedLowerBound = curLowerBound;
@@ -965,9 +976,7 @@ function assignUnconstrainedTypeVar(
     }
 
     if (constraints && !constraints.isLocked()) {
-        updateTypeVarType(
-            evaluator,
-            constraints,
+        constraints.setBounds(
             destType,
             newLowerBound,
             newUpperTypeBound,
@@ -994,7 +1003,7 @@ function assignConstrainedTypeVar(
 
     const curUpperBound = curEntry?.upperBound;
     const curLowerBound = curEntry?.lowerBound;
-    let forceRetainLiterals = false;
+    let retainLiterals = false;
 
     if (isTypeVar(srcType)) {
         if (
@@ -1129,7 +1138,7 @@ function assignConstrainedTypeVar(
         );
         return false;
     } else if (isLiteralTypeOrUnion(constrainedType)) {
-        forceRetainLiterals = true;
+        retainLiterals = true;
     }
 
     if (curLowerBound && !isAnyOrUnknown(curLowerBound)) {
@@ -1159,7 +1168,7 @@ function assignConstrainedTypeVar(
                 )
             ) {
                 if (constraints && !constraints.isLocked()) {
-                    updateTypeVarType(evaluator, constraints, destType, constrainedType, curUpperBound);
+                    constraints.setBounds(destType, constrainedType, curUpperBound);
                 }
             } else {
                 diag?.addMessage(
@@ -1174,7 +1183,7 @@ function assignConstrainedTypeVar(
     } else {
         // Assign the type to the type var.
         if (constraints && !constraints.isLocked()) {
-            updateTypeVarType(evaluator, constraints, destType, constrainedType, curUpperBound, forceRetainLiterals);
+            constraints.setBounds(destType, constrainedType, curUpperBound, retainLiterals);
         }
     }
 
@@ -1400,7 +1409,7 @@ function logTypeVarConstraintSet(evaluator: TypeEvaluator, context: ConstraintSe
 
     context.getTypeVars().forEach((entry) => {
         const typeVarName = `${indent}${entry.typeVar.shared.name}`;
-        const lowerBound = entry.lowerBoundNoLiterals ?? entry.lowerBound;
+        const lowerBound = entry.lowerBound;
         const upperBound = entry.upperBound;
 
         // Log the lower and upper bounds.

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

@@ -27,10 +27,12 @@ const maxConstraintSetCount = 1024;
 export interface TypeVarConstraints {
     typeVar: TypeVarType;
 
-    // Running constraints for the solved type variable as constraints are added.
+    // Bounds for solved type variable as constraints are added.
     lowerBound?: Type | undefined;
-    lowerBoundNoLiterals?: Type | undefined;
     upperBound?: Type | undefined;
+
+    // Should the lower bound include literal values?
+    retainLiterals?: boolean;
 }
 
 // Records the constraints information for a set of type variables
@@ -55,7 +57,7 @@ export class ConstraintSet {
         const constraintSet = new ConstraintSet();
 
         this._typeVarMap.forEach((value) => {
-            constraintSet.setBounds(value.typeVar, value.lowerBound, value.lowerBoundNoLiterals, value.upperBound);
+            constraintSet.setBounds(value.typeVar, value.lowerBound, value.upperBound, value.retainLiterals);
         });
 
         if (this._scopeIds) {
@@ -119,13 +121,13 @@ export class ConstraintSet {
         return score;
     }
 
-    setBounds(reference: TypeVarType, lowerBound: Type | undefined, lowerBoundNoLiterals?: Type, upperBound?: Type) {
-        const key = TypeVarType.getNameWithScope(reference);
+    setBounds(typeVar: TypeVarType, lowerBound: Type | undefined, upperBound?: Type, retainLiterals?: boolean) {
+        const key = TypeVarType.getNameWithScope(typeVar);
         this._typeVarMap.set(key, {
-            typeVar: reference,
+            typeVar,
             lowerBound,
-            lowerBoundNoLiterals,
             upperBound,
+            retainLiterals,
         });
     }
 
@@ -133,8 +135,8 @@ export class ConstraintSet {
         this._typeVarMap.forEach(cb);
     }
 
-    getTypeVar(reference: TypeVarType): TypeVarConstraints | undefined {
-        const key = TypeVarType.getNameWithScope(reference);
+    getTypeVar(typeVar: TypeVarType): TypeVarConstraints | undefined {
+        const key = TypeVarType.getNameWithScope(typeVar);
         return this._typeVarMap.get(key);
     }
 
@@ -220,6 +222,12 @@ export class ConstraintTracker {
         this._isLocked = clone._isLocked;
     }
 
+    copyBounds(entry: TypeVarConstraints) {
+        this._constraintSets.forEach((set) => {
+            set.setBounds(entry.typeVar, entry.lowerBound, entry.upperBound, entry.retainLiterals);
+        });
+    }
+
     // Copy the specified constraint sets into this type var context.
     addConstraintSets(contexts: ConstraintSet[]) {
         assert(contexts.length > 0);
@@ -258,11 +266,11 @@ export class ConstraintTracker {
         return this._constraintSets.every((set) => set.isEmpty());
     }
 
-    setBounds(reference: TypeVarType, lowerBound: Type | undefined, lowerBoundNoLiterals?: Type, upperBound?: Type) {
+    setBounds(typeVar: TypeVarType, lowerBound: Type | undefined, upperBound?: Type, retainLiterals?: boolean) {
         assert(!this._isLocked);
 
         return this._constraintSets.forEach((set) => {
-            set.setBounds(reference, lowerBound, lowerBoundNoLiterals, upperBound);
+            set.setBounds(typeVar, lowerBound, upperBound, retainLiterals);
         });
     }
 

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

@@ -758,12 +758,7 @@ function assignClassToProtocolInternal(
                 const typeArgEntry = protocolConstraints.getMainConstraintSet().getTypeVar(typeParam);
 
                 if (typeArgEntry) {
-                    destConstraints.setBounds(
-                        typeParam,
-                        typeArgEntry?.lowerBound,
-                        typeArgEntry?.lowerBoundNoLiterals,
-                        typeArgEntry?.upperBound
-                    );
+                    destConstraints.copyBounds(typeArgEntry);
                 }
             }
         }
@@ -786,7 +781,7 @@ function createProtocolConstraints(
         const entry = constraints?.getMainConstraintSet().getTypeVar(typeParam);
 
         if (entry) {
-            protocolConstraints.setBounds(typeParam, entry.lowerBound, entry.lowerBoundNoLiterals, entry.upperBound);
+            protocolConstraints.copyBounds(entry);
         } else if (destType.priv.typeArgs && index < destType.priv.typeArgs.length) {
             let typeArg = destType.priv.typeArgs[index];
             let flags: AssignTypeFlags;

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

@@ -101,11 +101,11 @@ import {
 } from './codeFlowTypes';
 import {
     addConstraintsForExpectedType,
+    applySourceSolutionToConstraints,
     applyUnificationVars,
     assignTypeVar,
     solveConstraintSet,
     solveConstraints,
-    updateTypeVarType,
 } from './constraintSolver';
 import { ConstraintSet, ConstraintTracker } from './constraintTracker';
 import { createFunctionFromConstructor, getBoundInitMethod, validateConstructorArgs } from './constructors';
@@ -279,7 +279,6 @@ import {
     addConditionToType,
     addTypeVarsToListIfUnique,
     applySolvedTypeVars,
-    applySourceContextTypeVars,
     areTypesSame,
     buildSolutionFromSpecializedClass,
     combineSameSizedTuples,
@@ -11617,7 +11616,10 @@ export function createTypeEvaluator(
                     // It's possible that one or more of the TypeVars or ParamSpecs
                     // in the constraints refer to TypeVars that were solved in
                     // the paramSpecConstraints. Apply these solved TypeVars accordingly.
-                    applySourceContextTypeVars(constraints, solveConstraints(evaluatorInterface, paramSpecConstraints));
+                    applySourceSolutionToConstraints(
+                        constraints,
+                        solveConstraints(evaluatorInterface, paramSpecConstraints)
+                    );
                 }
             });
         }
@@ -23075,13 +23077,11 @@ export function createTypeEvaluator(
                     typeArgType = i < srcTypeArgs.length ? srcTypeArgs[i] : UnknownType.create();
                 }
 
-                updateTypeVarType(
-                    evaluatorInterface,
-                    destConstraints,
+                destConstraints.setBounds(
                     typeParam,
                     variance !== Variance.Contravariant ? typeArgType : undefined,
                     variance !== Variance.Covariant ? typeArgType : undefined,
-                    /* forceRetainLiterals */ true
+                    /* retainLiterals */ true
                 );
             }
         }
@@ -25835,7 +25835,7 @@ export function createTypeEvaluator(
 
         // Apply any solved source TypeVars to the dest TypeVar solutions. This
         // allows for higher-order functions to accept generic callbacks.
-        applySourceContextTypeVars(destConstraints, solveConstraints(evaluatorInterface, srcConstraints));
+        applySourceSolutionToConstraints(destConstraints, solveConstraints(evaluatorInterface, srcConstraints));
 
         return canAssign;
     }

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

@@ -11,7 +11,7 @@ import { appendArray } from '../common/collectionUtils';
 import { assert } from '../common/debug';
 import { ArgumentNode, ParamCategory } from '../parser/parseNodes';
 import { ConstraintSolution, ConstraintSolutionSet } from './constraintSolution';
-import { ConstraintSet, ConstraintTracker } from './constraintTracker';
+import { ConstraintTracker } from './constraintTracker';
 import { DeclarationType } from './declaration';
 import { Symbol, SymbolFlags, SymbolTable } from './symbol';
 import { isEffectivelyClassVar, isTypedDictMemberAccessedThroughIndex } from './symbolUtils';
@@ -1515,29 +1515,6 @@ export function applySolvedTypeVars(type: Type, solution: ConstraintSolution, op
     return transformer.apply(type, 0);
 }
 
-// Applies solved TypeVars from one context to this context.
-export function applySourceContextTypeVars(destContext: ConstraintTracker, srcSolution: ConstraintSolution) {
-    if (srcSolution.isEmpty()) {
-        return;
-    }
-
-    destContext.doForEachConstraintSet((set) => {
-        applySourceContextTypeVarsToSignature(set, srcSolution);
-    });
-}
-
-export function applySourceContextTypeVarsToSignature(set: ConstraintSet, srcSolution: ConstraintSolution) {
-    set.getTypeVars().forEach((entry) => {
-        const newLowerBound = entry.lowerBound ? applySolvedTypeVars(entry.lowerBound, srcSolution) : undefined;
-        const newLowerBoundNoLiterals = entry.lowerBoundNoLiterals
-            ? applySolvedTypeVars(entry.lowerBoundNoLiterals, srcSolution)
-            : undefined;
-        const newUpperBound = entry.upperBound ? applySolvedTypeVars(entry.upperBound, srcSolution) : undefined;
-
-        set.setBounds(entry.typeVar, newLowerBound, newLowerBoundNoLiterals, newUpperBound);
-    });
-}
-
 // Validates that a default type associated with a TypeVar does not refer to
 // other TypeVars or ParamSpecs that are out of scope.
 export function validateTypeVarDefault(

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

@@ -0,0 +1,29 @@
+# This sample tests the case involving nested calls that each use
+# constrained TypeVars but one is a subset of the other.
+
+from os import PathLike
+from typing import Any, AnyStr, Literal, LiteralString, TypeVar, overload
+
+AnyOrLiteralStr = TypeVar("AnyOrLiteralStr", str, bytes, LiteralString)
+
+
+def abspath(path: PathLike[AnyStr] | AnyStr) -> AnyStr:
+    ...
+
+
+@overload
+def dirname(p: PathLike[AnyStr]) -> AnyStr:
+    ...
+
+
+@overload
+def dirname(p: AnyOrLiteralStr) -> AnyOrLiteralStr:
+    ...
+
+
+def dirname(p: Any) -> Any:
+    ...
+
+
+def func1(refpath: Literal["-"]):
+    reveal_type(dirname(abspath(refpath)), expected_text="str")

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

@@ -554,6 +554,12 @@ test('ConstrainedTypeVar19', () => {
     TestUtils.validateResults(analysisResults, 1);
 });
 
+test('ConstrainedTypeVar20', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['constrainedTypeVar20.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('MissingTypeArg1', () => {
     const configOptions = new ConfigOptions(Uri.empty());