MDS/openide
1
0
Fork 0
mirror of https://gitflic.ru/project/openide/openide.git synced 2026-04-20 13:31:28 +07:00
Code Issues Packages Projects Releases Wiki Activity

PY-76659 Make CFG traversal iterative for name resolution and type inference

Browse Source 
Swap back to ConcurrentMap and add more Registry checks

softValueMap doesn't play well with assumeType. Imagine this: you compute type A and store it, then compute type B that depends on A. But since softValueMap uses soft references, A might get garbage-collected. If that happens, you can override A with assumeType, and now B is out of sync — it's still based on the old version of A.

This kind of issue is unlikely to show up in small, artificial examples, but it could easily become a problem in real-world projects, especially large ones. It breaks the assumption that if a type is in the map, then everything it depends on should still be there too.

fix after rebase

Added nullability annotations to the AssumptionContext constructor parameters to improve type safety.

dm-checkpoint-id: 1VH4Od1GtvAo

Squash for easier rebase


Co-authored-by: Space Team <noreply@jetbrains.team>

Merge-request: IJ-MR-146970
Merged-by: Aleksandr Govenko <aleksandr.govenko@jetbrains.com>

GitOrigin-RevId: cf2fc232c2c5b35a037396b5f85be3129a5efd3a
This commit is contained in:
Aleksandr.Govenko
2025-04-20 16:03:28 +00:00
committed by intellij-monorepo-bot
parent 1bb2d6c5fc
commit a757828b06
7 changed files with 320 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

137
python/python-psi-api/src/com/jetbrains/python/psi/types/TypeEvalContext.java
View File

@@ -4,6 +4,8 @@ package com.jetbrains.python.psi.types;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.util.Pair;
import com.intellij.openapi.util.RecursionManager;
import com.intellij.openapi.util.Ref;
import com.intellij.openapi.util.registry.Registry;
import com.intellij.openapi.util.text.StringUtil;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
@@ -21,9 +23,11 @@ import org.jetbrains.annotations.Nullable;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
public final class TypeEvalContext {
public sealed class TypeEvalContext {
/**
* This class ensures that only {@link TypeEvalContext} instances can directly invoke
@@ -45,13 +49,30 @@ public final class TypeEvalContext {
private final ThreadLocal<ProcessingContext> myProcessingContext = ThreadLocal.withInitial(ProcessingContext::new);
private final Map<PyTypedElement, PyType> myEvaluated = CollectionFactory.createConcurrentSoftValueMap();
private final Map<PyCallable, PyType> myEvaluatedReturn = CollectionFactory.createConcurrentSoftValueMap();
protected final Map<PyTypedElement, PyType> myEvaluated = createMap();
protected final Map<PyCallable, PyType> myEvaluatedReturn = createMap();
/**
* AssumptionContext invariant requires that if type is in the map,
* it's dependencies are also in the map, so we can't use softValueMap.
* Temporary solution until we know assumeType works as expected.
* @see TypeEvalContext#assumeType(PyTypedElement, PyType, Function)
*/
private static <T> Map<T, PyType> createMap() {
if (Registry.is("python.use.better.control.flow.type.inference")) {
return new ConcurrentHashMap<>();
}
return CollectionFactory.createConcurrentSoftValueMap();
}
private TypeEvalContext(boolean allowDataFlow, boolean allowStubToAST, boolean allowCallContext, @Nullable PsiFile origin) {
myConstraints = new TypeEvalConstraints(allowDataFlow, allowStubToAST, allowCallContext, origin);
}
private TypeEvalContext(@NotNull TypeEvalConstraints constraints) {
myConstraints = constraints;
}
@Override
public String toString() {
return String.format("TypeEvalContext(%b, %b, %s)", myConstraints.myAllowDataFlow, myConstraints.myAllowStubToAST,
@@ -170,17 +191,55 @@ public final class TypeEvalContext {
return myTrace != null;
}
@ApiStatus.Internal
public <R> @Nullable R assumeType(@NotNull PyTypedElement element, @Nullable PyType type, @NotNull Function<TypeEvalContext, R> func) {
if (getKnownType(element) != null) {
// Temporary solution, as overwriting known type might introduce inconsistencies with its dependencies.
return null;
}
AssumptionContext context = new AssumptionContext(this, element, type);
R result = null;
try {
result = func.apply(context);
}
finally {
element.getManager().dropResolveCaches();
}
return result;
}
@ApiStatus.Internal
public boolean hasAssumptions() {
return this instanceof AssumptionContext;
}
protected @Nullable PyType getKnownType(final @NotNull PyTypedElement element) {
final PyType cachedType = myEvaluated.get(element);
if (cachedType != null) {
assertValid(cachedType, element);
return cachedType;
}
return null;
}
protected @Nullable PyType getKnownReturnType(final @NotNull PyCallable callable) {
final PyType cachedType = myEvaluatedReturn.get(callable);
if (cachedType != null) {
assertValid(cachedType, callable);
return cachedType;
}
return null;
}
public @Nullable PyType getType(final @NotNull PyTypedElement element) {
final PyType knownType = getKnownType(element);
if (knownType != null) {
return knownType == PyNullType.INSTANCE ? null : knownType;
}
return RecursionManager.doPreventingRecursion(
Pair.create(element, this),
false,
() -> {
PyType cachedType = myEvaluated.get(element);
if (cachedType != null) {
assertValid(cachedType, element);
return cachedType == PyNullType.INSTANCE ? null : cachedType;
}
PyType type = element.getType(this, Key.INSTANCE);
assertValid(type, element);
myEvaluated.put(element, type == null ? PyNullType.INSTANCE : type);
@@ -190,16 +249,15 @@ public final class TypeEvalContext {
}
public @Nullable PyType getReturnType(final @NotNull PyCallable callable) {
final PyType knownReturnType = getKnownReturnType(callable);
if (knownReturnType != null) {
return knownReturnType == PyNullType.INSTANCE ? null : knownReturnType;
}
return RecursionManager.doPreventingRecursion(
Pair.create(callable, this),
false,
() -> {
PyType cachedType = myEvaluatedReturn.get(callable);
if (cachedType != null) {
assertValid(cachedType, callable);
return cachedType == PyNullType.INSTANCE ? null : cachedType;
}
PyType type = callable.getReturnType(this, Key.INSTANCE);
final PyType type = callable.getReturnType(this, Key.INSTANCE);
assertValid(type, callable);
myEvaluatedReturn.put(callable, type == null ? PyNullType.INSTANCE : type);
return type;
@@ -306,4 +364,53 @@ public final class TypeEvalContext {
private static final PyNullType INSTANCE = new PyNullType();
}
private static final class AssumptionContext extends TypeEvalContext {
@NotNull final TypeEvalContext myParent;
private AssumptionContext(@NotNull TypeEvalContext parent, @NotNull PyTypedElement element, @Nullable PyType type) {
super(parent.myConstraints);
myParent = parent;
myEvaluated.put(element, type == null ? PyNullType.INSTANCE : type);
}
@Override
protected @Nullable PyType getKnownType(@NotNull PyTypedElement element) {
final PyType knownType = super.getKnownType(element);
if (knownType != null) {
return knownType;
}
return myParent.getKnownType(element);
}
@Override
protected @Nullable PyType getKnownReturnType(@NotNull PyCallable callable) {
final PyType knownReturnType = super.getKnownReturnType(callable);
if (knownReturnType != null) {
return knownReturnType;
}
return myParent.getKnownReturnType(callable);
}
@Override
public void trace(String message, Object... args) {
myParent.trace(message, args);
}
@Override
public void traceIndent() {
myParent.traceIndent();
}
@Override
public void traceUnindent() {
myParent.traceUnindent();
}
@Override
public boolean equals(Object o) {
// Otherwise, it can be equal to other AssumptionContext with same constraints
return this == o;
}
}
}

3
python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml
View File

@@ -464,6 +464,9 @@
<categoryKey>INTN.category.python</categoryKey>
</intentionAction>
<registryKey defaultValue="true"
description="Enable experimental better handling of recursive cases in type inference"
key="python.use.better.control.flow.type.inference"/>
<registryKey defaultValue="true"
description="Enable pyi stubs distributed with numpy (see https://github.com/numpy/numpy/issues/18565)"
key="enable.numpy.pyi.stubs"/>

6
python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyCallExpressionHelper.java
View File

@@ -156,9 +156,9 @@ public final class PyCallExpressionHelper {
call,
resolveContext,
it -> ContainerUtil.concat(
getExplicitResolveResults(call, it),
getImplicitResolveResults(call, it),
getRemoteResolveResults(call, it))
getExplicitResolveResults(call, it),
getImplicitResolveResults(call, it),
getRemoteResolveResults(call, it))
);
}

118
python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java
View File

@@ -2,11 +2,13 @@
package com.jetbrains.python.psi.impl;
import com.intellij.codeInsight.controlflow.ConditionalInstruction;
import com.intellij.codeInsight.controlflow.ControlFlowUtil;
import com.intellij.codeInsight.controlflow.Instruction;
import com.intellij.diagnostic.PluginException;
import com.intellij.lang.ASTNode;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.util.Ref;
import com.intellij.openapi.util.registry.Registry;
import com.intellij.psi.*;
import com.intellij.psi.impl.source.resolve.FileContextUtil;
import com.intellij.psi.util.PsiTreeUtil;
@@ -14,6 +16,7 @@ import com.intellij.psi.util.QualifiedName;
import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.PythonRuntimeService;
import com.jetbrains.python.codeInsight.controlflow.ControlFlowCache;
import com.jetbrains.python.codeInsight.controlflow.PyTypeAssertionEvaluator;
import com.jetbrains.python.codeInsight.controlflow.ReadWriteInstruction;
import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
@@ -41,6 +44,8 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
private static final Logger LOG = Logger.getInstance(PyReferenceExpressionImpl.class);
private static final int MAX_CFG_ITERATIONS = 30;
private volatile @Nullable QualifiedName myQualifiedName = null;
public PyReferenceExpressionImpl(@NotNull ASTNode astNode) {
@@ -465,6 +470,87 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
@NotNull TypeEvalContext context,
@NotNull PyExpression anchor,
@NotNull ScopeOwner scopeOwner) {
if (!Registry.is("python.use.better.control.flow.type.inference")) {
return getTypeByControlFlowOld(name, context, anchor, scopeOwner);
}
final PyAugAssignmentStatement augAssignment = PsiTreeUtil.getParentOfType(anchor, PyAugAssignmentStatement.class);
final PyElement element = augAssignment != null ? augAssignment : anchor;
final Instruction[] flow = ControlFlowCache.getControlFlow(scopeOwner).getInstructions();
final int thisInstructionIdx = ControlFlowUtil.findInstructionNumberByElement(flow, element);
if (thisInstructionIdx == -1) return null;
final Instruction thisInstruction = flow[thisInstructionIdx];
final List<Instruction> defs = PyDefUseUtil.getLatestDefs(scopeOwner, name, element, true, false, context);
// null means empty set of possible types, Ref(null) means Any
final @Nullable Ref<PyType> typeOfEarlierDefinitions = StreamEx.of(defs)
.filter(def -> def.num() < thisInstruction.num())
.map(def -> getTypeFromInstruction(context, anchor, def))
.nonNull()
.collect(PyTypeUtil.toUnionFromRef());
// If earlier definitions were not found, variable may be unbound. Choose Any as type.
PyType deducedType = Ref.deref(typeOfEarlierDefinitions);
final var laterDefs = StreamEx.of(defs).filter(def -> def.num() > thisInstruction.num()).toList();
if (laterDefs.isEmpty()) {
return deducedType;
}
for (int i = 0; i < MAX_CFG_ITERATIONS; i++) {
final @Nullable Ref<PyType> typeOfLaterDefinitions = context.assumeType(anchor, deducedType, (ctx) -> {
return StreamEx.of(laterDefs)
.map(def -> getTypeFromInstruction(ctx, anchor, def))
.nonNull()
.collect(PyTypeUtil.toUnionFromRef());
});
if (typeOfLaterDefinitions == null) {
return deducedType;
}
PyType newType = PyUnionType.union(deducedType, typeOfLaterDefinitions.get());
if (Objects.equals(deducedType, newType)) {
return deducedType;
}
deducedType = newType;
}
return deducedType;
}
private static @Nullable Ref<PyType> getTypeFromInstruction(@NotNull TypeEvalContext context,
@NotNull PyExpression anchor,
@NotNull Instruction instr) {
if (instr instanceof ReadWriteInstruction readWriteInstruction) {
return readWriteInstruction.getType(context, anchor);
}
if (instr instanceof ConditionalInstruction conditionalInstruction) {
final PyType conditionType = context.getType((PyTypedElement)conditionalInstruction.getCondition());
if (conditionType instanceof PyNarrowedType narrowedType && narrowedType.isBound()) {
var arguments = narrowedType.getOriginal().getArguments(null);
if (!arguments.isEmpty()) {
var firstArgument = arguments.get(0);
PyType type = narrowedType.getNarrowedType();
if (firstArgument instanceof PyReferenceExpression && type != null) {
@Nullable PyType initial = context.getType(firstArgument);
boolean positive = conditionalInstruction.getResult() ^ narrowedType.getNegated();
if (narrowedType.getTypeIs()) {
return PyTypeAssertionEvaluator.createAssertionType(initial, type, positive, context);
}
return Ref.create((positive) ? type : initial);
}
}
}
}
return null;
}
private static PyType getTypeByControlFlowOld(@NotNull String name,
@NotNull TypeEvalContext context,
@NotNull PyExpression anchor,
@NotNull ScopeOwner scopeOwner) {
final PyAugAssignmentStatement augAssignment = PsiTreeUtil.getParentOfType(anchor, PyAugAssignmentStatement.class);
final PyElement element = augAssignment != null ? augAssignment : anchor;
try {
@@ -476,22 +562,22 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
return readWriteInstruction.getType(context, anchor);
}
if (instr instanceof ConditionalInstruction conditionalInstruction) {
if (context.getType((PyTypedElement)conditionalInstruction.getCondition()) instanceof PyNarrowedType narrowedType
&& narrowedType.isBound()) {
var arguments = narrowedType.getOriginal().getArguments(null);
if (!arguments.isEmpty()) {
var firstArgument = arguments.get(0);
PyType type = narrowedType.getNarrowedType();
if (firstArgument instanceof PyReferenceExpression && type != null) {
@Nullable PyType initial = context.getType(firstArgument);
boolean positive = conditionalInstruction.getResult() ^ narrowedType.getNegated();
if (narrowedType.getTypeIs()) {
return PyTypeAssertionEvaluator.createAssertionType(initial, type, positive, context);
}
return Ref.create((positive) ? type : initial);
}
}
}
if (context.getType((PyTypedElement)conditionalInstruction.getCondition()) instanceof PyNarrowedType narrowedType
&& narrowedType.isBound()) {
var arguments = narrowedType.getOriginal().getArguments(null);
if (!arguments.isEmpty()) {
var firstArgument = arguments.get(0);
PyType type = narrowedType.getNarrowedType();
if (firstArgument instanceof PyReferenceExpression && type != null) {
@Nullable PyType initial = context.getType(firstArgument);
boolean positive = conditionalInstruction.getResult() ^ narrowedType.getNegated();
if (narrowedType.getTypeIs()) {
return PyTypeAssertionEvaluator.createAssertionType(initial, type, positive, context);
}
return Ref.create((positive) ? type : initial);
}
}
}
}
return null;
})

3
python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java
View File

@@ -88,7 +88,8 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference
public ResolveResult @NotNull [] multiResolve(final boolean incompleteCode) {
if (USE_CACHE) {
final ResolveCache cache = ResolveCache.getInstance(getElement().getProject());
return cache.resolveWithCaching(this, CachingResolver.INSTANCE, true, incompleteCode);
final boolean actuallyIncomplete = incompleteCode || myContext.getTypeEvalContext().hasAssumptions();
return cache.resolveWithCaching(this, CachingResolver.INSTANCE, true, actuallyIncomplete);
}
else {
return multiResolveInner();

21
python/python-psi-impl/src/com/jetbrains/python/refactoring/PyDefUseUtil.java
View File

@@ -22,6 +22,7 @@ import com.intellij.codeInsight.controlflow.Instruction;
import com.intellij.openapi.util.Comparing;
import com.intellij.openapi.util.Ref;
import com.intellij.openapi.util.Version;
import com.intellij.openapi.util.registry.Registry;
import com.intellij.psi.PsiElement;
import com.intellij.psi.util.QualifiedName;
import com.jetbrains.python.codeInsight.controlflow.*;
@@ -75,7 +76,6 @@ public final class PyDefUseUtil {
ControlFlowUtil.iteratePrev(startNum, instructions,
instruction -> {
if (instruction instanceof PyWithContextExitInstruction withExit) {
// probably should remove acceptTypeAssertions and make context nullable
if (!withExit.isSuppressingExceptions(context)) {
return ControlFlowUtil.Operation.CONTINUE;
}
@@ -86,8 +86,7 @@ public final class PyDefUseUtil {
result.add(typeGuardInstruction);
return ControlFlowUtil.Operation.CONTINUE;
}
// not a back edge
if (instruction.num() < startNum &&
if (isNotBackEdge(instruction.num(), startNum) &&
context.getOrigin() == callInstruction.getElement().getContainingFile()) {
var newContext = (MAX_CONTROL_FLOW_SIZE > instructions.length)
? TypeEvalContext.codeAnalysis(context.getOrigin().getProject(), context.getOrigin())
@@ -96,9 +95,8 @@ public final class PyDefUseUtil {
}
}
final PsiElement element = instruction.getElement();
if (acceptTypeAssertions
&& instruction instanceof ConditionalInstruction conditionalInstruction
&& instruction.num() < startNum) {
if (isNotBackEdge(instruction.num(), startNum)
&& acceptTypeAssertions && instruction instanceof ConditionalInstruction conditionalInstruction) {
if (conditionalInstruction.getCondition() instanceof PyTypedElement typedElement && context.getOrigin() == typedElement.getContainingFile()) {
var newContext = (MAX_CONTROL_FLOW_SIZE > instructions.length)
? TypeEvalContext.codeAnalysis(context.getOrigin().getProject(), context.getOrigin())
@@ -135,6 +133,17 @@ public final class PyDefUseUtil {
return new ArrayList<>(result);
}
/**
* New analysis handles back edges separately.
* @see com.jetbrains.python.psi.impl.PyReferenceExpressionImpl#getTypeByControlFlow(String, TypeEvalContext, PyExpression, ScopeOwner)
*/
private static boolean isNotBackEdge(int instNum, int startNum) {
if (Registry.is("python.use.better.control.flow.type.inference")) {
return true;
}
return instNum < startNum;
}
private static int findStartInstructionId(@NotNull PsiElement startAnchor, Instruction @NotNull [] instructions) {
PsiElement realCfgAnchor = startAnchor;
final PyAugAssignmentStatement augAssignment = PyAugAssignmentStatementNavigator.getStatementByTarget(startAnchor);

74
python/testSrc/com/jetbrains/python/Py3TypeTest.java
View File

@@ -104,6 +104,78 @@ public class Py3TypeTest extends PyTestCase {
""");
}
// PY-76659
public void testRecursiveResolve() {
doTest("int",
"""
x = 42
while x:
x = x + 1
expr = x""");
}
// PY-76659
public void testRecursiveResolve2() {
doTest("int",
"""
x = 42
b: bool = ...
while x:
if b:
x = x + 1
expr = x
else:
x = x - 1
""");
}
// PY-76659
public void testDeclareAfterUse() {
doTest("int | Any",
"""
from typing import Any, TypeGuard
def is_positive_integer(value: Any) -> TypeGuard[int]:
return isinstance(value, int) and value > 0
def bar() -> object:
return 321
def foo():
for i in range(1, 100):
if i > 1:
expr = x
x = bar()
if not is_positive_integer(x):
break
""");
}
// PY-76659
public void testClassChain() {
doTest("B | C | D | A",
"""
class A:
def bar() -> "B":
return B()
class B:
def bar() -> "C":
return C()
class C:
def bar() -> "D":
return D()
class D:
def bar() -> A:
return A()
def foo(b):
x = A()
while b:
x = x.bar()
expr = x""");
}
// PY-6702
public void testYieldFromType() {
doTest("str | int | float",
@@ -788,7 +860,7 @@ public class Py3TypeTest extends PyTestCase {
async for expr in a:
print(expr)""");
}
// PY-60714
public void testAsyncIteratorUnwrapsCoroutineFromAnext() {
doTest("bytes", """