From 7a81f4dfba5446418890409b694e30e75f07953d Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Sat, 18 Mar 2017 21:59:01 +0300
Subject: [PATCH] PY-22087 Don't mark comprehension variables used in f-strings
 as unused

I needed to tweak anchor element (that has the corresponding instruction
in CFG) to flow-sensitively resolve names used in f-strings inside
comprehensions.
---
 .../documentation/doctest/PyDocReference.java     | 15 ++++++++++-----
 .../PyUnusedLocalInspectionVisitor.java           |  3 ++-
 .../test.py                                       | 10 ++++++++++
 .../jetbrains/python/PythonInspectionsTest.java   |  5 +++++
 4 files changed, 27 insertions(+), 6 deletions(-)
 create mode 100644 python/testData/inspections/UnusedLocalFStringReferencesInComprehensions/test.py

diff --git a/python/src/com/jetbrains/python/documentation/doctest/PyDocReference.java b/python/src/com/jetbrains/python/documentation/doctest/PyDocReference.java
index a8875d190658..35cb5a81a710 100644
--- a/python/src/com/jetbrains/python/documentation/doctest/PyDocReference.java
+++ b/python/src/com/jetbrains/python/documentation/doctest/PyDocReference.java
@@ -99,15 +99,20 @@ public class PyDocReference extends PyReferenceImpl {
   @Nullable
   private PsiElement getScopeControlFlowAnchor(@NotNull PsiLanguageInjectionHost host) {
     if (isInsideFormattedStringNode(host)) {
-      final PsiElement comprehensionPart = PsiTreeUtil.findFirstParent(host, PyDocReference::isComprehensionResultOrComponent);
-      if (comprehensionPart != null) {
-        return comprehensionPart;
-      }
-      return PsiTreeUtil.getParentOfType(host, PyStatement.class);
+      return getControlFlowAnchorForFString((PyStringLiteralExpression)host);
     }
     return null;
   }
 
+  @Nullable
+  public static PsiElement getControlFlowAnchorForFString(@NotNull PyStringLiteralExpression host) {
+    final PsiElement comprehensionPart = PsiTreeUtil.findFirstParent(host, PyDocReference::isComprehensionResultOrComponent);
+    if (comprehensionPart != null) {
+      return comprehensionPart;
+    }
+    return PsiTreeUtil.getParentOfType(host, PyStatement.class);
+  }
+
   private static boolean isComprehensionResultOrComponent(@NotNull PsiElement element) {
     // Any comprehension component and its result are represented as children expressions of the comprehension element.
     // Only they have respective nodes in CFG and thus can be used as anchors for getResultsFromProcessor().
diff --git a/python/src/com/jetbrains/python/inspections/PyUnusedLocalInspectionVisitor.java b/python/src/com/jetbrains/python/inspections/PyUnusedLocalInspectionVisitor.java
index 1bf9294a98a6..548f62edfe52 100644
--- a/python/src/com/jetbrains/python/inspections/PyUnusedLocalInspectionVisitor.java
+++ b/python/src/com/jetbrains/python/inspections/PyUnusedLocalInspectionVisitor.java
@@ -34,6 +34,7 @@ import com.jetbrains.python.codeInsight.controlflow.ReadWriteInstruction;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.codeInsight.dataflow.scope.Scope;
 import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
+import com.jetbrains.python.documentation.doctest.PyDocReference;
 import com.jetbrains.python.inspections.quickfix.AddFieldQuickFix;
 import com.jetbrains.python.inspections.quickfix.PyRemoveParameterQuickFix;
 import com.jetbrains.python.inspections.quickfix.PyRemoveStatementQuickFix;
@@ -104,7 +105,7 @@ public class PyUnusedLocalInspectionVisitor extends PyInspectionVisitor {
   public void visitPyStringLiteralExpression(PyStringLiteralExpression pyString) {
     final ScopeOwner owner = ScopeUtil.getScopeOwner(pyString);
     if (owner != null && !(owner instanceof PsiFile)) {
-      final PyStatement instrAnchor = PsiTreeUtil.getParentOfType(pyString, PyStatement.class);
+      final PsiElement instrAnchor = PyDocReference.getControlFlowAnchorForFString(pyString);
       if (instrAnchor == null) return;
       final Instruction[] instructions = ControlFlowCache.getControlFlow(owner).getInstructions();
       final int startInstruction = ControlFlowUtil.findInstructionNumberByElement(instructions, instrAnchor);
diff --git a/python/testData/inspections/UnusedLocalFStringReferencesInComprehensions/test.py b/python/testData/inspections/UnusedLocalFStringReferencesInComprehensions/test.py
new file mode 100644
index 000000000000..56aaa0a661d6
--- /dev/null
+++ b/python/testData/inspections/UnusedLocalFStringReferencesInComprehensions/test.py
@@ -0,0 +1,10 @@
+def f():
+    [f'{x1}' for x1 in range(10)]
+    [[f'{x2}' for _ in range(10)] for x2 in range(10)]
+    [[42 for _ in range(10) if f'{x3}'] for x3 in range(10)]
+
+    [42 for x4 in range(10) if f'{x4}']
+    [42 for x5 in range(10) if [f'{x5}' for _ in range(10)]]
+    [42 for x6 in range(10) if [42 for _ in range(10) if f'{x6}']]
+
+    [f'{x7}' for _ in range(10) for x7 in range(10)]
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PythonInspectionsTest.java b/python/testSrc/com/jetbrains/python/PythonInspectionsTest.java
index 04d7f86dc1e7..5edefbc07b05 100644
--- a/python/testSrc/com/jetbrains/python/PythonInspectionsTest.java
+++ b/python/testSrc/com/jetbrains/python/PythonInspectionsTest.java
@@ -123,6 +123,11 @@ public class PythonInspectionsTest extends PyTestCase {
     doHighlightingTest(PyUnusedLocalInspection.class, LanguageLevel.PYTHON36);
   }
 
+  // PY-22087
+  public void testUnusedLocalFStringReferencesInComprehensions() {
+    doHighlightingTest(PyUnusedLocalInspection.class, LanguageLevel.PYTHON36);
+  }
+
   // PY-8219
   public void testUnusedLocalDoctestReference() {
     doHighlightingTest(PyUnusedLocalInspection.class);