From 2e8150a2e97b77410f0052386ee83a7c567d5f4b Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Fri, 16 Aug 2013 14:07:24 +0400
Subject: [PATCH] Added stubs for docstrings of target expressions

---
 .../python/psi/PyTargetExpression.java        |  3 +-
 .../psi/stubs/PyTargetExpressionStub.java     |  3 +
 .../python/documentation/DocStringUtil.java   | 18 -----
 .../documentation/PyDocumentationBuilder.java |  2 +-
 .../python/psi/PyFileElementType.java         |  3 +-
 .../psi/impl/PyTargetExpressionImpl.java      | 76 +++++++++++++------
 .../stubs/PyTargetExpressionElementType.java  | 16 +++-
 .../stubs/PyTargetExpressionStubImpl.java     | 16 +++-
 .../psi/stubs/PyInstanceAttributeIndex.java   |  5 ++
 .../python/psi/stubs/PyVariableNameIndex.java |  5 ++
 .../stubs/TargetExpressionDocString.py        |  6 ++
 .../com/jetbrains/python/PyStubsTest.java     |  8 ++
 12 files changed, 111 insertions(+), 50 deletions(-)
 create mode 100644 python/testData/stubs/TargetExpressionDocString.py

diff --git a/python/psi-api/src/com/jetbrains/python/psi/PyTargetExpression.java b/python/psi-api/src/com/jetbrains/python/psi/PyTargetExpression.java
index 930a4c310af9..d27d84191bc7 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/PyTargetExpression.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/PyTargetExpression.java
@@ -12,7 +12,8 @@ import org.jetbrains.annotations.Nullable;
 /**
  * @author yole
  */
-public interface PyTargetExpression extends PyQualifiedExpression, PsiNamedElement, PsiNameIdentifierOwner, StubBasedPsiElement<PyTargetExpressionStub> {
+public interface PyTargetExpression extends PyQualifiedExpression, PsiNamedElement, PsiNameIdentifierOwner, PyDocStringOwner,
+                                            StubBasedPsiElement<PyTargetExpressionStub> {
   PyTargetExpression[] EMPTY_ARRAY = new PyTargetExpression[0];
 
   /**
diff --git a/python/psi-api/src/com/jetbrains/python/psi/stubs/PyTargetExpressionStub.java b/python/psi-api/src/com/jetbrains/python/psi/stubs/PyTargetExpressionStub.java
index aedd7d3740e0..7ec0c2e7788e 100644
--- a/python/psi-api/src/com/jetbrains/python/psi/stubs/PyTargetExpressionStub.java
+++ b/python/psi-api/src/com/jetbrains/python/psi/stubs/PyTargetExpressionStub.java
@@ -45,4 +45,7 @@ public interface PyTargetExpressionStub extends NamedStub<PyTargetExpression> {
 
   @Nullable
   <T extends CustomTargetExpressionStub> T getCustomStub(Class<T> stubClass);
+
+  @Nullable
+  String getDocString();
 }
diff --git a/python/src/com/jetbrains/python/documentation/DocStringUtil.java b/python/src/com/jetbrains/python/documentation/DocStringUtil.java
index 5f25e1e2b82a..812f518963e6 100644
--- a/python/src/com/jetbrains/python/documentation/DocStringUtil.java
+++ b/python/src/com/jetbrains/python/documentation/DocStringUtil.java
@@ -73,24 +73,6 @@ public class DocStringUtil {
     return false;
   }
 
-  @Nullable
-  public static PyStringLiteralExpression getAttributeDocString(@NotNull PyTargetExpression attr) {
-    if (attr.getParent() instanceof PyAssignmentStatement) {
-      final PyAssignmentStatement assignment = (PyAssignmentStatement)attr.getParent();
-      PsiElement nextSibling = assignment.getNextSibling();
-      while (nextSibling != null && (nextSibling instanceof PsiWhiteSpace || nextSibling instanceof PsiComment)) {
-        nextSibling = nextSibling.getNextSibling();
-      }
-      if (nextSibling instanceof PyExpressionStatement) {
-        final PyExpression expression = ((PyExpressionStatement)nextSibling).getExpression();
-        if (expression instanceof PyStringLiteralExpression) {
-          return (PyStringLiteralExpression)expression;
-        }
-      }
-    }
-    return null;
-  }
-
   @Nullable
   public static String getAttributeDocComment(@NotNull PyTargetExpression attr) {
     if (attr.getParent() instanceof PyAssignmentStatement) {
diff --git a/python/src/com/jetbrains/python/documentation/PyDocumentationBuilder.java b/python/src/com/jetbrains/python/documentation/PyDocumentationBuilder.java
index a54cc473d5d5..281f7b186865 100644
--- a/python/src/com/jetbrains/python/documentation/PyDocumentationBuilder.java
+++ b/python/src/com/jetbrains/python/documentation/PyDocumentationBuilder.java
@@ -446,7 +446,7 @@ class PyDocumentationBuilder {
       .addItem(" of class ").addWith(PythonDocumentationProvider.LinkMyClass, $().addWith(TagCode, $(cls.getName()))).addItem(BR)
     ;
 
-    final String docString = PyPsiUtils.strValue(DocStringUtil.getAttributeDocString((PyTargetExpression)myElement));
+    final String docString = ((PyTargetExpression)myElement).getDocStringValue();
     if (docString != null) {
       addFormattedDocString(myElement, docString, myBody, myEpilog);
     }
diff --git a/python/src/com/jetbrains/python/psi/PyFileElementType.java b/python/src/com/jetbrains/python/psi/PyFileElementType.java
index 355818dede2e..14154f11d6a9 100644
--- a/python/src/com/jetbrains/python/psi/PyFileElementType.java
+++ b/python/src/com/jetbrains/python/psi/PyFileElementType.java
@@ -47,7 +47,8 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
 
   @Override
   public int getStubVersion() {
-    return 47;
+    // Don't forget to update versions of indexes that use the updated stub-based elements
+    return 48;
   }
 
   @Nullable
diff --git a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
index 804c9c1255e2..97f0a39be9f2 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
@@ -4,10 +4,7 @@ import com.intellij.lang.ASTNode;
 import com.intellij.openapi.extensions.Extensions;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.Ref;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiPolyVariantReference;
-import com.intellij.psi.PsiReference;
-import com.intellij.psi.ResolveResult;
+import com.intellij.psi.*;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.search.LocalSearchScope;
 import com.intellij.psi.search.SearchScope;
@@ -25,7 +22,6 @@ 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.DocStringUtil;
-import com.jetbrains.python.psi.StructuredDocString;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.references.PyQualifiedReference;
 import com.jetbrains.python.psi.impl.references.PyTargetReference;
@@ -120,10 +116,14 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
       if (pyType != null) {
         return pyType;
       }
+      PyType type = getTypeFromDocString();
+      if (type != null) {
+        return type;
+      }
       if (!context.maySwitchToAST(this)) {
         final PsiElement value = getStub() != null ? findAssignedValueByStub(context) : findAssignedValue();
         if (value instanceof PyTypedElement) {
-          final PyType type = context.getType((PyTypedElement)value);
+          type = context.getType((PyTypedElement)value);
           if (type instanceof PyNoneType) {
             return null;
           }
@@ -135,10 +135,6 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
         }
         return null;
       }
-      PyType type = getTypeFromDocString(this);
-      if (type != null) {
-        return type;
-      }
       type = getTypeFromComment(this);
       if (type != null) {
         return type;
@@ -249,29 +245,27 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
   }
 
   @Nullable
-  public static PyType getTypeFromDocString(PyTargetExpressionImpl targetExpression) {
+  private PyType getTypeFromDocString() {
     String typeName = null;
-    final String docString = PyPsiUtils.strValue(DocStringUtil.getAttributeDocString(targetExpression));
-    if (docString != null) {
-      StructuredDocString targetDocString = DocStringUtil.parse(docString);
-      if (targetDocString != null) {
-        typeName = targetDocString.getParamType(null);
-        if (typeName == null) {
-          typeName = targetDocString.getParamType(targetExpression.getName());
-        }
+    final String name = getName();
+    final StructuredDocString targetDocString = getStructuredDocString();
+    if (targetDocString != null) {
+      typeName = targetDocString.getParamType(null);
+      if (typeName == null) {
+        typeName = targetDocString.getParamType(name);
       }
     }
-    if (typeName == null && PyUtil.isAttribute(targetExpression)) {
-      final PyClass cls = targetExpression.getContainingClass();
+    if (typeName == null && PyUtil.isAttribute(this)) {
+      final PyClass cls = getContainingClass();
       if (cls != null) {
         final StructuredDocString classDocString = cls.getStructuredDocString();
         if (classDocString != null) {
-          typeName = classDocString.getParamType(targetExpression.getName());
+          typeName = classDocString.getParamType(name);
         }
       }
     }
     if (typeName != null) {
-      return PyTypeParser.getTypeByName(targetExpression, typeName);
+      return PyTypeParser.getTypeByName(this, typeName);
     }
     return null;
   }
@@ -613,4 +607,40 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
     }
     return super.getElementLocation();
   }
+
+  @Nullable
+  @Override
+  public String getDocStringValue() {
+    final PyTargetExpressionStub stub = getStub();
+    if (stub != null) {
+      return stub.getDocString();
+    }
+    return DocStringUtil.getDocStringValue(this);
+  }
+
+  @Nullable
+  @Override
+  public StructuredDocString getStructuredDocString() {
+    return DocStringUtil.getStructuredDocString(this);
+  }
+
+  @Nullable
+  @Override
+  public PyStringLiteralExpression getDocStringExpression() {
+    final PsiElement parent = getParent();
+    if (parent instanceof PyAssignmentStatement) {
+      final PyAssignmentStatement assignment = (PyAssignmentStatement)parent;
+      PsiElement nextSibling = assignment.getNextSibling();
+      while (nextSibling != null && (nextSibling instanceof PsiWhiteSpace || nextSibling instanceof PsiComment)) {
+        nextSibling = nextSibling.getNextSibling();
+      }
+      if (nextSibling instanceof PyExpressionStatement) {
+        final PyExpression expression = ((PyExpressionStatement)nextSibling).getExpression();
+        if (expression instanceof PyStringLiteralExpression) {
+          return (PyStringLiteralExpression)expression;
+        }
+      }
+    }
+    return null;
+  }
 }
diff --git a/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java b/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java
index d7f2669adeeb..205211190414 100644
--- a/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java
+++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java
@@ -12,6 +12,7 @@ import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.io.StringRef;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyNames;
+import com.jetbrains.python.documentation.DocStringUtil;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyQualifiedName;
 import com.jetbrains.python.psi.impl.PyTargetExpressionImpl;
@@ -48,10 +49,11 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
   public PyTargetExpressionStub createStub(@NotNull final PyTargetExpression psi, final StubElement parentStub) {
     final String name = psi.getName();
     final PyExpression assignedValue = psi.findAssignedValue();
+    final String docString = DocStringUtil.getDocStringValue(psi);
     for (CustomTargetExpressionStubType customStubType : getCustomStubTypes()) {
       CustomTargetExpressionStub customStub = customStubType.createStub(psi);
       if (customStub != null) {
-        return new PyTargetExpressionStubImpl(name, customStub, parentStub);
+        return new PyTargetExpressionStubImpl(name, docString, customStub, parentStub);
       }
     }
     PyTargetExpressionStub.InitializerType initializerType = PyTargetExpressionStub.InitializerType.Other;
@@ -67,12 +69,14 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
         initializer = ((PyReferenceExpression) callee).asQualifiedName();
       }
     }
-    return new PyTargetExpressionStubImpl(name, initializerType, initializer, psi.getQualifier() != null, parentStub);
+    return new PyTargetExpressionStubImpl(name, docString, initializerType, initializer, psi.getQualifier() != null, parentStub);
   }
 
   public void serialize(@NotNull final PyTargetExpressionStub stub, @NotNull final StubOutputStream stream)
       throws IOException {
     stream.writeName(stub.getName());
+    final String docString = stub.getDocString();
+    stream.writeUTFFast(docString != null ? docString : "");
     stream.writeVarInt(stub.getInitializerType().getIndex());
     final CustomTargetExpressionStub customStub = stub.getCustomStub(CustomTargetExpressionStub.class);
     if (customStub != null) {
@@ -89,20 +93,24 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
   public PyTargetExpressionStub deserialize(@NotNull final StubInputStream stream, final StubElement parentStub)
       throws IOException {
     String name = StringRef.toString(stream.readName());
+    String docString = stream.readUTFFast();
+    if (docString.isEmpty()) {
+      docString = null;
+    }
     PyTargetExpressionStub.InitializerType initializerType = PyTargetExpressionStub.InitializerType.fromIndex(stream.readVarInt());
     if (initializerType == PyTargetExpressionStub.InitializerType.Custom) {
       final String typeName = stream.readName().getString();
       for(CustomTargetExpressionStubType type: getCustomStubTypes()) {
         if (type.getClass().getCanonicalName().equals(typeName)) {
           CustomTargetExpressionStub stub = type.deserializeStub(stream);
-          return new PyTargetExpressionStubImpl(name, stub, parentStub);
+          return new PyTargetExpressionStubImpl(name, docString, stub, parentStub);
         }
       }
       throw new IOException("Unknown custom stub type " + typeName);
     }
     PyQualifiedName initializer = PyQualifiedName.deserialize(stream);
     boolean isQualified = stream.readBoolean();
-    return new PyTargetExpressionStubImpl(name, initializerType, initializer, isQualified, parentStub);
+    return new PyTargetExpressionStubImpl(name, docString, initializerType, initializer, isQualified, parentStub);
   }
 
   public boolean shouldCreateStub(final ASTNode node) {
diff --git a/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionStubImpl.java b/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionStubImpl.java
index eb7b2338ce89..3519a71da5c3 100644
--- a/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionStubImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionStubImpl.java
@@ -16,19 +16,24 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
   private final InitializerType myInitializerType;
   private final PyQualifiedName myInitializer;
   private final boolean myQualified;
+  @Nullable private final String myDocString;
 
   private final CustomTargetExpressionStub myCustomStub;
 
-  public PyTargetExpressionStubImpl(String name, CustomTargetExpressionStub customStub, StubElement parent) {
+  public PyTargetExpressionStubImpl(String name,
+                                    @Nullable String docString,
+                                    CustomTargetExpressionStub customStub,
+                                    StubElement parent) {
     super(parent, PyElementTypes.TARGET_EXPRESSION);
     myName = name;
     myInitializerType = InitializerType.Custom;
     myInitializer = null;
     myQualified = false;
     myCustomStub = customStub;
+    myDocString = docString;
   }
   
-  public PyTargetExpressionStubImpl(final String name, final InitializerType initializerType,
+  public PyTargetExpressionStubImpl(final String name, @Nullable String docString, final InitializerType initializerType,
                                     final PyQualifiedName initializer,
                                     final boolean qualified,
                                     final StubElement parentStub) {
@@ -39,6 +44,7 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
     myInitializer = initializer;
     myQualified = qualified;
     myCustomStub = null;
+    myDocString = docString;
   }
 
   public String getName() {
@@ -67,6 +73,12 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
     return null;
   }
 
+  @Nullable
+  @Override
+  public String getDocString() {
+    return myDocString;
+  }
+
   @Override
   public String toString() {
     return "PyTargetExpressionStub(name=" + myName + ")";
diff --git a/python/src/com/jetbrains/python/psi/stubs/PyInstanceAttributeIndex.java b/python/src/com/jetbrains/python/psi/stubs/PyInstanceAttributeIndex.java
index 408cf5183794..57967b31a39d 100644
--- a/python/src/com/jetbrains/python/psi/stubs/PyInstanceAttributeIndex.java
+++ b/python/src/com/jetbrains/python/psi/stubs/PyInstanceAttributeIndex.java
@@ -16,6 +16,11 @@ import java.util.Collection;
 public class PyInstanceAttributeIndex extends StringStubIndexExtension<PyTargetExpression> {
   public static final StubIndexKey<String, PyTargetExpression> KEY = StubIndexKey.createIndexKey("Py.instanceAttribute.name");
 
+  @Override
+  public int getVersion() {
+    return super.getVersion() + 1;
+  }
+
   @NotNull
   @Override
   public StubIndexKey<String, PyTargetExpression> getKey() {
diff --git a/python/src/com/jetbrains/python/psi/stubs/PyVariableNameIndex.java b/python/src/com/jetbrains/python/psi/stubs/PyVariableNameIndex.java
index 392e974ef3cd..162734c51e6d 100644
--- a/python/src/com/jetbrains/python/psi/stubs/PyVariableNameIndex.java
+++ b/python/src/com/jetbrains/python/psi/stubs/PyVariableNameIndex.java
@@ -16,6 +16,11 @@ import java.util.Collection;
 public class PyVariableNameIndex extends StringStubIndexExtension<PyTargetExpression> {
   public static final StubIndexKey<String, PyTargetExpression> KEY = StubIndexKey.createIndexKey("Py.variable.shortName");
 
+  @Override
+  public int getVersion() {
+    return super.getVersion() + 1;
+  }
+
   @NotNull
   public StubIndexKey<String, PyTargetExpression> getKey() {
     return KEY;
diff --git a/python/testData/stubs/TargetExpressionDocString.py b/python/testData/stubs/TargetExpressionDocString.py
new file mode 100644
index 000000000000..e6051a8b74e7
--- /dev/null
+++ b/python/testData/stubs/TargetExpressionDocString.py
@@ -0,0 +1,6 @@
+class C(object):
+    foo = None
+    """Foo docstring."""
+
+    bar = None
+
diff --git a/python/testSrc/com/jetbrains/python/PyStubsTest.java b/python/testSrc/com/jetbrains/python/PyStubsTest.java
index fdf34339ee4c..95306b724a08 100644
--- a/python/testSrc/com/jetbrains/python/PyStubsTest.java
+++ b/python/testSrc/com/jetbrains/python/PyStubsTest.java
@@ -402,6 +402,14 @@ public class PyStubsTest extends PyTestCase {
     }
     classes = PyClassNameIndex.find("Foo", project, GlobalSearchScope.allScope(project));
     assertEquals(classes.size(), 1);
+  }
 
+  public void testTargetExpressionDocString() {
+    final PyFile file = getTestFile();
+    final PyClass c = file.findTopLevelClass("C");
+    assertNotNull(c);
+    final PyTargetExpression foo = c.findClassAttribute("foo", false);
+    final String docString = foo.getDocStringValue();
+    assertEquals("Foo docstring.", docString);
   }
 }