Efficient doc string annotator; tests.

python/src/com/jetbrains/python/validation/DocStringAnnotator.java

@@ -17,67 +17,42 @@
 package com.jetbrains.python.validation;
 
 import com.intellij.lang.annotation.Annotation;
-import com.intellij.psi.PsiComment;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiWhiteSpace;
 import com.jetbrains.python.PyHighlighter;
-import com.jetbrains.python.PyTokenTypes;
-import com.jetbrains.python.psi.*;
-import com.jetbrains.python.psi.impl.PyFileImpl;
-import org.jetbrains.annotations.Nullable;
+import com.jetbrains.python.PythonDosStringFinder;
+import com.jetbrains.python.psi.PyClass;
+import com.jetbrains.python.psi.PyFile;
+import com.jetbrains.python.psi.PyFunction;
+import com.jetbrains.python.psi.PyStringLiteralExpression;
 
 /**
- * Created by IntelliJ IDEA.
- * User: yole
- * Date: 13.06.2005
- * Time: 22:42:04
- * To change this template use File | Settings | File Templates.
+ * Highlights doc strings in classes, functions, and files.
  */
 public class DocStringAnnotator extends PyAnnotator {
+
   @Override
-  public void visitPyExpressionStatement(final PyExpressionStatement node) {
-    PsiElement parent = node.getParent();
-    if (parent instanceof PyFileImpl && isFirstNotCommentChild(node, parent)) {
-      annotateDocStringStmt(node);
-    }
-    else if (parent instanceof PyStatementList && isFirstNotCommentChild(node, parent)) {
-      PsiElement stmtParent = parent.getParent();
-      if (stmtParent instanceof PyFunction || stmtParent instanceof PyClass) {
-        annotateDocStringStmt(node);
-      }
-    }
+  public void visitPyFile(PyFile node) {
+    annotateDocStringStmt(PythonDosStringFinder.find(node));
   }
 
-  private static boolean isFirstNotCommentChild(final PsiElement node, final PsiElement parent) {
-    PsiElement[] children = parent.getChildren();
-    for(PsiElement child: children) {
-      if (child == node) return true;
-      if (!(child instanceof PsiComment) && !(child instanceof PsiWhiteSpace)) break;
-    }
-    return false;
+
+  @Override
+  public void visitPyFunction(PyFunction node) {
+    annotateDocStringStmt(PythonDosStringFinder.find(node.getStatementList()));
   }
 
-  private void annotateDocStringStmt(PyExpressionStatement stmt) {
-    final PyStringLiteralExpression docString = statementAsDocString(stmt);
-    if (docString != null) {
-      Annotation ann = getHolder().createInfoAnnotation(docString, null);
+  @Override
+  public void visitPyClass(PyClass node) {
+    annotateDocStringStmt(PythonDosStringFinder.find(node.getStatementList()));
+  }
+
+  private void annotateDocStringStmt(PyStringLiteralExpression stmt) {
+    if (stmt != null) {
+      Annotation ann = getHolder().createInfoAnnotation(stmt, null);
       ann.setTextAttributes(PyHighlighter.PY_DOC_COMMENT);
     }
   }
 
-  @Nullable
-  public static PyStringLiteralExpression statementAsDocString(final PyExpressionStatement stmt) {
-    PyExpression expr = stmt.getExpression();
-    if (expr instanceof PyLiteralExpression) {
-      PyLiteralExpression litExpr = (PyLiteralExpression)expr;
-      PsiElement elt = litExpr.getFirstChild();
-      if (elt != null && elt.getNode().getElementType() == PyTokenTypes.STRING_LITERAL) {
-        return (PyStringLiteralExpression) litExpr;
-      }
-    }
-    return null;
-  }
-
+  /*
   @Nullable
   public static String findDocString(final PyStatementList statementList) {
     final PyStatement[] statements = statementList.getStatements();
@@ -88,4 +63,5 @@ public class DocStringAnnotator extends PyAnnotator {
     }
     return null;
   }
+  */
 }

python/testData/highlighting/docStrings.py

@@ -0,0 +1,16 @@
+# bg is always black.
+# effect is white
+# doc comment: blue bold
+def <info descr="null" type="INFORMATION">foo</info>():
+  <info descr="null" type="INFORMATION" foreground="0x0000ff" background="0x000000" effectcolor="0xffffff" effecttype="BOXED" fonttype="1">"Func doc string"</info>
+  pass
+
+class <info descr="null" type="INFORMATION">Boo</info>:
+  <info descr="null" type="INFORMATION" foreground="0x0000ff" background="0x000000" effectcolor="0xffffff" effecttype="BOXED" fonttype="1">"Class doc string"</info>
+  pass
+
+class <info descr="null" type="INFORMATION">Moo</info>:
+  def <info descr="null" type="INFORMATION">meth</info>(self):
+    <info descr="null" type="INFORMATION" foreground="0x0000ff" background="0x000000" effectcolor="0xffffff" effecttype="BOXED" fonttype="1">"Meth doc string"</info>
+    pass
+ 

python/testData/highlighting/malformedStringEscaped.py

@@ -1,2 +1,3 @@
 # missing last quote, escaped
-<error descr="Missing closing quote [']">'abc\'</error>
+a = 1 # make below not a doc comment
+<error descr="Missing closing quote [']">'abc\'</error>

python/testData/highlighting/malformedStringUnterminated.py

@@ -1,2 +1,3 @@
 # missing last quote, plain
-<error descr="Missing closing quote [']">'abc</error>
+a = 1 # to make the text below non-doc
+<error descr="Missing closing quote [']">'abc</error>

python/testSrc/com/jetbrains/python/PythonHighlightingTest.java

@@ -11,6 +11,7 @@ import com.intellij.openapi.editor.markup.EffectType;
 import java.awt.*;
 
 /**
+ * Test highlighting added by annotators.
  * @author yole
  */
 public class PythonHighlightingTest extends DaemonAnalyzerTestCase {
@@ -39,6 +40,19 @@ public class PythonHighlightingTest extends DaemonAnalyzerTestCase {
     doTest();
   }
 
+  public void testDocStrings() throws Exception {
+    EditorColorsManager manager = EditorColorsManager.getInstance();
+    EditorColorsScheme scheme = (EditorColorsScheme)manager.getGlobalScheme().clone();
+    manager.addColorsScheme(scheme);
+    EditorColorsManager.getInstance().setGlobalScheme(scheme);
+
+    TextAttributesKey xKey = TextAttributesKey.find("PY.DOC_COMMENT");
+    TextAttributes xAttributes = new TextAttributes(Color.blue, Color.black, Color.white, EffectType.BOXED, Font.BOLD);
+    scheme.setAttributes(xKey, xAttributes);
+
+    doTest();
+  }
+
   public void testReturnOutsideOfFunction() throws Exception {
     doTest();
   }