From 8a4dfe563c60baaa233b7e064d336e2b0bcc4ff7 Mon Sep 17 00:00:00 2001 From: Vladimir Koshelev Date: Tue, 9 Jul 2024 15:57:37 +0200 Subject: [PATCH] [pycharm] add type eval context to PyDefUseUtil GitOrigin-RevId: a6973f48d9e1f5c473209df25ff507be482e8904 --- .../python/psi/impl/PyReferenceExpressionImpl.java | 2 +- .../python/psi/impl/references/PyReferenceImpl.java | 4 ++-- .../jetbrains/python/refactoring/PyDefUseUtil.java | 12 +++++++----- .../refactoring/inline/PyInlineLocalHandler.java | 5 +++-- 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java index dc61ef1c3193..418cfca9997e 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java @@ -477,7 +477,7 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere final PyAugAssignmentStatement augAssignment = PsiTreeUtil.getParentOfType(anchor, PyAugAssignmentStatement.class); final PyElement element = augAssignment != null ? augAssignment : anchor; try { - final List defs = PyDefUseUtil.getLatestDefs(scopeOwner, name, element, true, false); + final List defs = PyDefUseUtil.getLatestDefs(scopeOwner, name, element, true, false, context); // null means empty set of possible types, Ref(null) means Any final @Nullable Ref combinedType = StreamEx.of(defs) .select(ReadWriteInstruction.class) diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java index 07dad59e8dcd..163140b46e97 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java +++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java @@ -278,7 +278,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference if (outermostNestedClass != null) { final List instructions = - PyDefUseUtil.getLatestDefs(resolvedOwner, referencedName, outermostNestedClass, false, true); + PyDefUseUtil.getLatestDefs(resolvedOwner, referencedName, outermostNestedClass, false, true, typeEvalContext); return resolveToLatestDefs(instructions, outermostNestedClass, referencedName, typeEvalContext); } @@ -334,7 +334,7 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference protected List getLatestDefinitions(@NotNull String referencedName, @Nullable ScopeOwner resolvedOwner, @Nullable PsiElement realContext) { - return PyDefUseUtil.getLatestDefs(resolvedOwner, referencedName, realContext, false, true); + return PyDefUseUtil.getLatestDefs(resolvedOwner, referencedName, realContext, false, true, myContext.getTypeEvalContext()); } private boolean allInOwnScopeComprehensions(@NotNull Collection elements) { diff --git a/python/python-psi-impl/src/com/jetbrains/python/refactoring/PyDefUseUtil.java b/python/python-psi-impl/src/com/jetbrains/python/refactoring/PyDefUseUtil.java index e52bf83c6557..37c41cf2351f 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/refactoring/PyDefUseUtil.java +++ b/python/python-psi-impl/src/com/jetbrains/python/refactoring/PyDefUseUtil.java @@ -28,6 +28,7 @@ import com.jetbrains.python.codeInsight.controlflow.ScopeOwner; import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyAugAssignmentStatementNavigator; +import com.jetbrains.python.psi.types.TypeEvalContext; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; @@ -42,14 +43,14 @@ public final class PyDefUseUtil { @NotNull public static List getLatestDefs(ScopeOwner block, String varName, PsiElement anchor, boolean acceptTypeAssertions, - boolean acceptImplicitImports) { - return getLatestDefs(ControlFlowCache.getControlFlow(block), varName, anchor, acceptTypeAssertions, acceptImplicitImports); + boolean acceptImplicitImports, @NotNull TypeEvalContext context) { + return getLatestDefs(ControlFlowCache.getControlFlow(block), varName, anchor, acceptTypeAssertions, acceptImplicitImports, context); } @NotNull public static List getLatestDefs(ControlFlow controlFlow, String varName, PsiElement anchor, boolean acceptTypeAssertions, - boolean acceptImplicitImports) { + boolean acceptImplicitImports, @NotNull TypeEvalContext context) { final Instruction[] instructions = controlFlow.getInstructions(); final PyAugAssignmentStatement augAssignment = PyAugAssignmentStatementNavigator.getStatementByTarget(anchor); if (augAssignment != null) { @@ -65,12 +66,13 @@ public final class PyDefUseUtil { instr = pred.iterator().next().num(); } } - final Collection result = getLatestDefs(varName, instructions, instr, acceptTypeAssertions, acceptImplicitImports); + final Collection result = getLatestDefs(varName, instructions, instr, acceptTypeAssertions, acceptImplicitImports, context); return new ArrayList<>(result); } private static Collection getLatestDefs(final String varName, final Instruction[] instructions, final int instr, - final boolean acceptTypeAssertions, final boolean acceptImplicitImports) { + final boolean acceptTypeAssertions, final boolean acceptImplicitImports, + @NotNull final TypeEvalContext context) { final Collection result = new LinkedHashSet<>(); ControlFlowUtil.iteratePrev(instr, instructions, instruction -> { diff --git a/python/python-psi-impl/src/com/jetbrains/python/refactoring/inline/PyInlineLocalHandler.java b/python/python-psi-impl/src/com/jetbrains/python/refactoring/inline/PyInlineLocalHandler.java index 8fc3c00b1ea4..0ae0a8b4b12c 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/refactoring/inline/PyInlineLocalHandler.java +++ b/python/python-psi-impl/src/com/jetbrains/python/refactoring/inline/PyInlineLocalHandler.java @@ -34,6 +34,7 @@ import com.jetbrains.python.PythonLanguage; import com.jetbrains.python.codeInsight.controlflow.ScopeOwner; import com.jetbrains.python.psi.*; import com.jetbrains.python.psi.impl.PyPsiUtils; +import com.jetbrains.python.psi.types.TypeEvalContext; import com.jetbrains.python.refactoring.PyDefUseUtil; import com.jetbrains.python.refactoring.PyReplaceExpressionUtil; import org.jetbrains.annotations.Nls; @@ -277,7 +278,7 @@ public final class PyInlineLocalHandler extends InlineActionHandler { for (final PsiElement ref : refsToInline) { final List elems = new ArrayList<>(); - final List latestDefs = PyDefUseUtil.getLatestDefs(containerBlock, local.getName(), ref, false, false); + final List latestDefs = PyDefUseUtil.getLatestDefs(containerBlock, local.getName(), ref, false, false, TypeEvalContext.codeInsightFallback(project)); for (Instruction i : latestDefs) { elems.add(i.getElement()); } @@ -385,7 +386,7 @@ public final class PyInlineLocalHandler extends InlineActionHandler { PyTargetExpression local, Project project) { if (expr != null) { try { - final List candidates = PyDefUseUtil.getLatestDefs(containerBlock, local.getName(), expr, true, true); + final List candidates = PyDefUseUtil.getLatestDefs(containerBlock, local.getName(), expr, true, true, TypeEvalContext.codeInsightFallback(project)); if (candidates.size() == 1) { final PyStatement expression = getAssignmentByLeftPart((PyElement)candidates.get(0).getElement()); return Pair.create(expression, false);