From 7af9174046f024c4a8065c09d21fe8b1c675427f Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Fri, 8 May 2015 20:09:38 +0300
Subject: [PATCH] Inject Python into function annotations only if they look
 like type hints (PY-15810)

Currently PEP 484 is still a draft and PyCharm doesn't implement
ignore comments and decorators for suppressing type checking of
annotations. We shouldn't inject Python into annotations that are not
supposed to be type hints.
---
 .../PyTypingAnnotationInjector.java           | 15 ++++-
 .../com/jetbrains/python/PyTypingTest.java    | 64 +++++++++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/python/src/com/jetbrains/python/codeInsight/PyTypingAnnotationInjector.java b/python/src/com/jetbrains/python/codeInsight/PyTypingAnnotationInjector.java
index 16a7b86261b3..b460e0f03de9 100644
--- a/python/src/com/jetbrains/python/codeInsight/PyTypingAnnotationInjector.java
+++ b/python/src/com/jetbrains/python/codeInsight/PyTypingAnnotationInjector.java
@@ -24,16 +24,27 @@ import com.jetbrains.python.psi.PyStringLiteralExpression;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.regex.Pattern;
+
 /**
  * @author vlan
  */
 public class PyTypingAnnotationInjector extends PyInjectorBase {
+  public static final Pattern RE_TYPING_ANNOTATION = Pattern.compile("\\s*\\S+(\\[.*\\])?\\s*");
+
   @Nullable
   @Override
   public Language getInjectedLanguage(@NotNull PsiElement context) {
-    if (context instanceof PyStringLiteralExpression && PsiTreeUtil.getParentOfType(context, PyAnnotation.class, true) != null) {
-      return PyDocstringLanguageDialect.getInstance();
+    if (context instanceof PyStringLiteralExpression) {
+      final PyStringLiteralExpression expr = (PyStringLiteralExpression)context;
+      if (PsiTreeUtil.getParentOfType(context, PyAnnotation.class, true) != null && isTypingAnnotation(expr.getStringValue())) {
+        return PyDocstringLanguageDialect.getInstance();
+      }
     }
     return null;
   }
+
+  private static boolean isTypingAnnotation(@NotNull String s) {
+    return RE_TYPING_ANNOTATION.matcher(s).matches();
+  }
 }
diff --git a/python/testSrc/com/jetbrains/python/PyTypingTest.java b/python/testSrc/com/jetbrains/python/PyTypingTest.java
index 6ac1fdfd5c2e..b6126e9abeb6 100644
--- a/python/testSrc/com/jetbrains/python/PyTypingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypingTest.java
@@ -15,6 +15,15 @@
  */
 package com.jetbrains.python;
 
+import com.intellij.lang.injection.InjectedLanguageManager;
+import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.util.Pair;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiDocumentManager;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.PsiLanguageInjectionHost;
 import com.intellij.testFramework.LightProjectDescriptor;
 import com.jetbrains.python.documentation.PythonDocumentationProvider;
 import com.jetbrains.python.fixtures.PyTestCase;
@@ -25,6 +34,8 @@ import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.List;
+
 /**
  * Tests for a type system based on mypy's typing module.
  *
@@ -306,6 +317,59 @@ public class PyTypingTest extends PyTestCase {
            "        pass\n");
   }
 
+  public void testStringLiteralInjection() {
+    doTestInjectedText("class C:\n" +
+                       "    def foo(self, expr: '<caret>C'):\n" +
+                       "        pass\n",
+                       "C");
+  }
+
+  public void testStringLiteralInjectionParameterizedType() {
+    doTestInjectedText("from typing import Union, List\n" +
+                       "\n" +
+                       "class C:\n" +
+                       "    def foo(self, expr: '<caret>Union[List[C], C]'):\n" +
+                       "        pass\n",
+                       "Union[List[C], C]");
+  }
+
+  // PY-15810
+  public void testNoStringLiteralInjectionForNonTypingStrings() {
+    doTestNoInjectedText("class C:\n" +
+                         "    def foo(self, expr: '<caret>foo bar'):\n" +
+                         "        pass\n");
+  }
+
+  private void doTestNoInjectedText(@NotNull String text) {
+    myFixture.configureByText(PythonFileType.INSTANCE, text);
+    final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myFixture.getProject());
+    final PsiLanguageInjectionHost host = languageManager.getInjectionHost(getElementAtCaret());
+    assertNull(host);
+  }
+
+  private void doTestInjectedText(@NotNull String text, @NotNull String expected) {
+    myFixture.configureByText(PythonFileType.INSTANCE, text);
+    final InjectedLanguageManager languageManager = InjectedLanguageManager.getInstance(myFixture.getProject());
+    final PsiLanguageInjectionHost host = languageManager.getInjectionHost(getElementAtCaret());
+    assertNotNull(host);
+    final List<Pair<PsiElement, TextRange>> files = languageManager.getInjectedPsiFiles(host);
+    assertNotNull(files);
+    assertFalse(files.isEmpty());
+    final PsiElement injected = files.get(0).getFirst();
+    assertEquals(expected, injected.getText());
+  }
+
+  @NotNull
+  private PsiElement getElementAtCaret() {
+    final Editor editor = myFixture.getEditor();
+    final Document document = editor.getDocument();
+    final PsiFile file = PsiDocumentManager.getInstance(myFixture.getProject()).getPsiFile(document);
+    assertNotNull(file);
+    final PsiElement element = file.findElementAt(myFixture.getCaretOffset());
+    assertNotNull(element);
+    return element;
+  }
+
   private void doTest(@NotNull String expectedType, @NotNull String text) {
     myFixture.copyDirectoryToProject("typing", "");
     myFixture.configureByText(PythonFileType.INSTANCE, text);