diff --git a/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java b/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java index 0abdfaa9ff89..57d78d0a598e 100644 --- a/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java +++ b/python/src/com/jetbrains/python/psi/impl/references/PyReferenceImpl.java @@ -183,6 +183,8 @@ public class PyReferenceImpl implements PsiReferenceEx, PsiPolyVariantReference } private static boolean isInnerComprehension(PsiElement referenceElement, PsiElement definition) { + if (definition != null && definition.getParent() instanceof PyAssignmentExpression) return false; + final PyComprehensionElement definitionComprehension = PsiTreeUtil.getParentOfType(definition, PyComprehensionElement.class); if (definitionComprehension != null && PyUtil.isOwnScopeComprehension(definitionComprehension)) { final PyComprehensionElement elementComprehension = PsiTreeUtil.getParentOfType(referenceElement, PyComprehensionElement.class); diff --git a/python/testData/resolve/AssignmentExpressionsAndOuterVar.py b/python/testData/resolve/AssignmentExpressionsAndOuterVar.py new file mode 100644 index 000000000000..78f942b69039 --- /dev/null +++ b/python/testData/resolve/AssignmentExpressionsAndOuterVar.py @@ -0,0 +1,4 @@ +total = 0 +partial_sums = [total := total + v for v in values] +print("Total:", total) + \ No newline at end of file diff --git a/python/testSrc/com/jetbrains/python/PyResolveTest.java b/python/testSrc/com/jetbrains/python/PyResolveTest.java index e3b4231c0345..e6abb79e408c 100644 --- a/python/testSrc/com/jetbrains/python/PyResolveTest.java +++ b/python/testSrc/com/jetbrains/python/PyResolveTest.java @@ -1,6 +1,7 @@ // Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file. package com.jetbrains.python; +import com.intellij.openapi.util.Ref; import com.intellij.psi.*; import com.intellij.psi.util.PsiTreeUtil; import com.intellij.testFramework.UsefulTestCase; @@ -16,6 +17,8 @@ import com.jetbrains.python.psi.resolve.PyResolveContext; import com.jetbrains.python.psi.types.PyClassTypeImpl; import com.jetbrains.python.psi.types.TypeEvalContext; import com.jetbrains.python.pyi.PyiUtil; +import org.intellij.lang.annotations.Language; +import org.jetbrains.annotations.NotNull; public class PyResolveTest extends PyResolveTestCase { @Override @@ -41,6 +44,22 @@ public class PyResolveTest extends PyResolveTestCase { return ((PsiPolyVariantReference)ref).multiResolve(false); } + private T assertResolvesTo(@Language("TEXT") @NotNull String text, + @NotNull Class cls, + @NotNull String name) { + final Ref result = new Ref<>(); + + runWithLanguageLevel( + LanguageLevel.getLatest(), + () -> { + myFixture.configureByText(PythonFileType.INSTANCE, text); + result.set(assertResolveResult(PyResolveTestCase.findReferenceByMarker(myFixture.getFile()).resolve(), cls, name)); + } + ); + + return result.get(); + } + public void testClass() { assertResolvesTo(PyClass.class, "Test"); } @@ -1350,4 +1369,86 @@ public class PyResolveTest extends PyResolveTestCase { assertNotNull(function); assertFalse(PyiUtil.isOverload(function, context)); } + + // PY-33886 + public void testAssignmentExpressions() { + assertResolvesTo( + "if a := b:\n" + + " print(a)\n" + + " ", + PyTargetExpression.class, + "a" + ); + + assertResolvesTo( + "[y := 2, y**2]\n" + + " ", + PyTargetExpression.class, + "y" + ); + + assertResolvesTo( + "[y for x in data if (y := f(x))]\n" + + " ", + PyTargetExpression.class, + "y" + ); + + assertResolvesTo( + "len(lines := [])\n" + + "print(lines)\n" + + " ", + PyTargetExpression.class, + "lines" + ); + } + + // PY-33886 + public void testAssignmentExpressionGoesToOuterScope() { + assertResolvesTo( + "if any({(comment := line).startswith('#') for line in lines}):\n" + + " print(\"First comment:\", comment)\n" + + " ", + PyTargetExpression.class, + "comment" + ); + + assertResolvesTo( + "if all((nonblank := line).strip() == '' for line in lines):\n" + + " pass\n" + + "else:\n" + + " print(\"First non-blank line:\", nonblank)\n" + + " ", + PyTargetExpression.class, + "nonblank" + ); + } + + // PY-33886 + public void testAssignmentExpressionGoesFromOuterScope() { + assertResolvesTo( + "[[x * y for x in range(5)] for i in range(5) if (y := i)]\n" + + " ", + PyTargetExpression.class, + "y" + ); + } + + // PY-33886 + public void testAssignmentExpressionsAndOuterVar() { + runWithLanguageLevel( + LanguageLevel.PYTHON38, + () -> { + final ResolveResult[] results = multiResolve(); + + final PsiElement first = results[0].getElement(); + assertResolveResult(first, PyTargetExpression.class, "total"); + assertInstanceOf(first.getParent(), PyAssignmentStatement.class); + + final PsiElement second = results[1].getElement(); + assertResolveResult(second, PyTargetExpression.class, "total"); + assertInstanceOf(second.getParent(), PyAssignmentExpression.class); + } + ); + } }