From fc3d1b7beb6b40394ccd1ed82ef76806bde1ca50 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Sun, 21 Jan 2018 20:26:56 +0300
Subject: [PATCH] PY-16412 Don't suggest the intention for targets that doesn't
 allow annotations

---
 .../PyAnnotateVariableTypeIntention.java      | 18 ++++++++--
 .../notSuggestedForComprehensionTarget.py     |  1 +
 .../notSuggestedForGlobalTarget.py            |  3 ++
 .../notSuggestedForImportTarget.py            |  1 +
 ...signmentTargetWithAheadOfTimeAnnotation.py |  4 +++
 ...dForLocalAssignmentTargetWithAnnotation.py |  3 ++
 ...ForLocalAssignmentTargetWithTypeComment.py |  3 ++
 .../notSuggestedForNonlocalTarget.py          |  3 ++
 .../PyAnnotateVariableTypeIntentionTest.java  | 33 +++++++++++++++++++
 9 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForComprehensionTarget.py
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForGlobalTarget.py
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForImportTarget.py
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithAheadOfTimeAnnotation.py
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithAnnotation.py
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithTypeComment.py
 create mode 100644 python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForNonlocalTarget.py

diff --git a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java
index 1ad69906e3f9..49c1a8372285 100644
--- a/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java
+++ b/python/src/com/jetbrains/python/codeInsight/intentions/PyAnnotateVariableTypeIntention.java
@@ -63,10 +63,22 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction {
     return StreamEx.of(PyUtil.multiResolveTopPriority(elementAtCaret, resolveContext))
       .select(PyTargetExpression.class)
       .filter(target -> !index.isInLibraryClasses(target.getContainingFile().getVirtualFile()))
+      .filter(target -> canBeAnnotated(target))
       .filter(target -> !isAnnotated(target, typeEvalContext))
       .toList();
   }
 
+  private static boolean canBeAnnotated(@NotNull PyTargetExpression target) {
+    final PsiElement directParent = target.getParent();
+    if (directParent instanceof PyImportElement ||
+        directParent instanceof PyComprehensionForComponent ||
+        directParent instanceof PyGlobalStatement ||
+        directParent instanceof PyNonlocalStatement) {
+      return false;
+    }
+    return PsiTreeUtil.getParentOfType(target, PyWithItem.class, PyAssignmentStatement.class, PyForPart.class) != null;
+  }
+
   private static boolean isAnnotated(@NotNull PyTargetExpression target, @NotNull TypeEvalContext context) {
     // TODO filter out fields explicitly annotated as Any
     return new PyTypingTypeProvider().getReferenceType(target, context, null) != null;
@@ -105,9 +117,9 @@ public class PyAnnotateVariableTypeIntention extends PyBaseIntentionAction {
   private static String generateNestedTypeHint(@NotNull PyTargetExpression target) {
     final TypeEvalContext context = TypeEvalContext.userInitiated(target.getProject(), target.getContainingFile());
     final StringBuilder builder = new StringBuilder();
-    final PyElement validTargetContainer = PsiTreeUtil.getParentOfType(target, PyForPart.class, PyWithItem.class, PyAssignmentStatement.class);
-    assert validTargetContainer != null;
-    final PsiElement topmostTarget = PsiTreeUtil.findPrevParent(validTargetContainer, target);
+    final PyElement validTargetParent = PsiTreeUtil.getParentOfType(target, PyForPart.class, PyWithItem.class, PyAssignmentStatement.class);
+    assert validTargetParent != null;
+    final PsiElement topmostTarget = PsiTreeUtil.findPrevParent(validTargetParent, target);
     generateNestedTypeHint(topmostTarget, context, builder);
     return builder.toString();
   }
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForComprehensionTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForComprehensionTarget.py
new file mode 100644
index 000000000000..386d4f11dffa
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForComprehensionTarget.py
@@ -0,0 +1 @@
+[v<caret>ar for var in range(10)]
\ No newline at end of file
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForGlobalTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForGlobalTarget.py
new file mode 100644
index 000000000000..292e52a70ce0
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForGlobalTarget.py
@@ -0,0 +1,3 @@
+def func():
+    global var
+    v<caret>ar
\ No newline at end of file
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForImportTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForImportTarget.py
new file mode 100644
index 000000000000..fa08f698686c
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForImportTarget.py
@@ -0,0 +1 @@
+from sys import path as v<caret>ar
\ No newline at end of file
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithAheadOfTimeAnnotation.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithAheadOfTimeAnnotation.py
new file mode 100644
index 000000000000..6a63049eb153
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithAheadOfTimeAnnotation.py
@@ -0,0 +1,4 @@
+def func(x):
+    var: int
+    var = x
+    v<caret>ar
\ No newline at end of file
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithAnnotation.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithAnnotation.py
new file mode 100644
index 000000000000..9a79c94c9804
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithAnnotation.py
@@ -0,0 +1,3 @@
+def func(x):
+    var: int = x
+    v<caret>ar
\ No newline at end of file
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithTypeComment.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithTypeComment.py
new file mode 100644
index 000000000000..45e01ca94dce
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForLocalAssignmentTargetWithTypeComment.py
@@ -0,0 +1,3 @@
+def func(x):
+    var = x  # type: int
+    v<caret>ar
\ No newline at end of file
diff --git a/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForNonlocalTarget.py b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForNonlocalTarget.py
new file mode 100644
index 000000000000..7ea0550e9c49
--- /dev/null
+++ b/python/testData/intentions/PyAnnotateVariableTypeIntentionTest/notSuggestedForNonlocalTarget.py
@@ -0,0 +1,3 @@
+def func():
+    nonlocal var
+    v<caret>ar
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java
index e3bbbfeaac3d..7f090cb19a51 100644
--- a/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java
+++ b/python/testSrc/com/jetbrains/python/intentions/PyAnnotateVariableTypeIntentionTest.java
@@ -9,6 +9,35 @@ import org.jetbrains.annotations.NotNull;
  * @author Mikhail Golubev
  */
 public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase {
+
+  public void testNotSuggestedForLocalAssignmentTargetWithAnnotation() {
+    doNegativeTest();
+  }
+
+  public void testNotSuggestedForLocalAssignmentTargetWithAheadOfTimeAnnotation() {
+    doNegativeTest();
+  }
+
+  public void testNotSuggestedForLocalAssignmentTargetWithTypeComment() {
+    doNegativeTest();
+  }
+
+  public void testNotSuggestedForComprehensionTarget() {
+    doNegativeTest();
+  }
+
+  public void testNotSuggestedForImportTarget() {
+    doNegativeTest();
+  }
+
+  public void testNotSuggestedForGlobalTarget() {
+    doNegativeTest();
+  }
+
+  public void testNotSuggestedForNonlocalTarget() {
+    doNegativeTest();
+  }
+
   public void testAnnotationLocalSimpleAssignmentTarget() {
     doTestAnnotation();
   }
@@ -68,4 +97,8 @@ public class PyAnnotateVariableTypeIntentionTest extends PyIntentionTestCase {
   private void doTest(@NotNull LanguageLevel languageLevel) {
     doTest(PyBundle.message("INTN.annotate.types"), languageLevel);
   }
+
+  private void doNegativeTest() {
+    runWithLanguageLevel(LanguageLevel.PYTHON36, () -> doNegativeTest(PyBundle.message("INTN.annotate.types")));
+  }
 }