From 3dc0c0db148a3d294df33deada3c5b8a3cedb518 Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Wed, 10 Mar 2010 21:17:30 +0300
Subject: [PATCH] store initializer for target expressions in stubs

---
 .../com/jetbrains/python/psi/PyFileElementType.java |  2 +-
 .../jetbrains/python/psi/PyTargetExpression.java    |  9 ++++-----
 .../python/psi/impl/PyTargetExpressionImpl.java     |  9 +++++++++
 .../impl/stubs/PyTargetExpressionElementType.java   | 13 +++++++++++--
 .../psi/impl/stubs/PyTargetExpressionStubImpl.java  |  9 ++++++++-
 .../python/psi/stubs/PyTargetExpressionStub.java    |  2 ++
 python/testData/stubs/StubStructure.py              |  2 +-
 .../testSrc/com/jetbrains/python/PyStubsTest.java   |  1 +
 8 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/PyFileElementType.java b/python/src/com/jetbrains/python/psi/PyFileElementType.java
index 1cf598fb22b2..5e2bb1d76d01 100644
--- a/python/src/com/jetbrains/python/psi/PyFileElementType.java
+++ b/python/src/com/jetbrains/python/psi/PyFileElementType.java
@@ -24,7 +24,7 @@ public class PyFileElementType extends IStubFileElementType {
 
   @Override
   public int getStubVersion() {
-    return 19;
+    return 20;
   }
 
   @Override
diff --git a/python/src/com/jetbrains/python/psi/PyTargetExpression.java b/python/src/com/jetbrains/python/psi/PyTargetExpression.java
index 53bc4343e170..a4d9fd6a0f78 100644
--- a/python/src/com/jetbrains/python/psi/PyTargetExpression.java
+++ b/python/src/com/jetbrains/python/psi/PyTargetExpression.java
@@ -2,6 +2,7 @@ package com.jetbrains.python.psi;
 
 import com.intellij.psi.PsiNamedElement;
 import com.intellij.psi.StubBasedPsiElement;
+import com.jetbrains.python.psi.impl.PyQualifiedName;
 import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
 import org.jetbrains.annotations.Nullable;
 
@@ -11,11 +12,6 @@ import org.jetbrains.annotations.Nullable;
 public interface PyTargetExpression extends PyQualifiedExpression, PsiNamedElement, StubBasedPsiElement<PyTargetExpressionStub> {
   PyTargetExpression[] EMPTY_ARRAY = new PyTargetExpression[0];
 
-  /*
-  @Nullable
-  PyExpression getQualifier();
-  */
-
   /**
    * Find the value that maps to this target expression in an enclosing assignment expression.
    * Does not work with other expressions (e.g. if the target is in a 'for' loop).
@@ -24,4 +20,7 @@ public interface PyTargetExpression extends PyQualifiedExpression, PsiNamedEleme
    */
   @Nullable
   PyExpression findAssignedValue();
+
+  @Nullable
+  PyQualifiedName getAssignedQName();
 }
diff --git a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
index db3872cdbeae..dd87b563de7d 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
@@ -130,4 +130,13 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
     }
     return null;
   }
+
+  public PyQualifiedName getAssignedQName() {
+    final PyTargetExpressionStub stub = getStub();
+    if (stub != null) {
+      return stub.getInitializer();
+    }
+    final PyExpression value = findAssignedValue();
+    return value instanceof PyReferenceExpression ? ((PyReferenceExpression) value).asQualifiedName() : 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 c49afeb28278..33e55d04cfb0 100644
--- a/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java
+++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java
@@ -8,8 +8,11 @@ import com.intellij.psi.stubs.StubInputStream;
 import com.intellij.psi.stubs.StubOutputStream;
 import com.intellij.util.io.StringRef;
 import com.jetbrains.python.PyElementTypes;
+import com.jetbrains.python.psi.PyExpression;
+import com.jetbrains.python.psi.PyReferenceExpression;
 import com.jetbrains.python.psi.PyStubElementType;
 import com.jetbrains.python.psi.PyTargetExpression;
+import com.jetbrains.python.psi.impl.PyQualifiedName;
 import com.jetbrains.python.psi.impl.PyTargetExpressionImpl;
 import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
 
@@ -32,18 +35,24 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
   }
 
   public PyTargetExpressionStub createStub(final PyTargetExpression psi, final StubElement parentStub) {
-    return new PyTargetExpressionStubImpl(psi.getName(), parentStub);
+    final PyExpression assignedValue = psi.findAssignedValue();
+    final PyQualifiedName initializer = assignedValue instanceof PyReferenceExpression
+                                        ? ((PyReferenceExpression) assignedValue).asQualifiedName()
+                                        : null;
+    return new PyTargetExpressionStubImpl(psi.getName(), initializer, parentStub);
   }
 
   public void serialize(final PyTargetExpressionStub stub, final StubOutputStream dataStream)
       throws IOException {
     dataStream.writeName(stub.getName());
+    PyQualifiedName.serialize(stub.getInitializer(), dataStream);
   }
 
   public PyTargetExpressionStub deserialize(final StubInputStream dataStream, final StubElement parentStub)
       throws IOException {
     String name = StringRef.toString(dataStream.readName());
-    return new PyTargetExpressionStubImpl(name, parentStub);
+    PyQualifiedName initializer = PyQualifiedName.deserialize(dataStream);
+    return new PyTargetExpressionStubImpl(name, initializer, 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 4b8809fd656a..17085fdfc7ca 100644
--- a/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionStubImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionStubImpl.java
@@ -4,6 +4,7 @@ import com.intellij.psi.stubs.StubBase;
 import com.intellij.psi.stubs.StubElement;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.psi.PyTargetExpression;
+import com.jetbrains.python.psi.impl.PyQualifiedName;
 import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
 
 /**
@@ -11,13 +12,19 @@ import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
  */
 public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> implements PyTargetExpressionStub {
   private final String myName;
+  private final PyQualifiedName myInitializer;
 
-  public PyTargetExpressionStubImpl(final String name, final StubElement parentStub) {
+  public PyTargetExpressionStubImpl(final String name, final PyQualifiedName initializer, final StubElement parentStub) {
     super(parentStub, PyElementTypes.TARGET_EXPRESSION);
     myName = name;
+    myInitializer = initializer;
   }
 
   public String getName() {
     return myName;
   }
+
+  public PyQualifiedName getInitializer() {
+    return myInitializer;
+  }
 }
diff --git a/python/src/com/jetbrains/python/psi/stubs/PyTargetExpressionStub.java b/python/src/com/jetbrains/python/psi/stubs/PyTargetExpressionStub.java
index ceeb8cd9cb81..58147848c2f6 100644
--- a/python/src/com/jetbrains/python/psi/stubs/PyTargetExpressionStub.java
+++ b/python/src/com/jetbrains/python/psi/stubs/PyTargetExpressionStub.java
@@ -2,9 +2,11 @@ package com.jetbrains.python.psi.stubs;
 
 import com.intellij.psi.stubs.NamedStub;
 import com.jetbrains.python.psi.PyTargetExpression;
+import com.jetbrains.python.psi.impl.PyQualifiedName;
 
 /**
  * @author yole
  */
 public interface PyTargetExpressionStub extends NamedStub<PyTargetExpression> {
+  PyQualifiedName getInitializer();
 }
diff --git a/python/testData/stubs/StubStructure.py b/python/testData/stubs/StubStructure.py
index 3e67a8f06e21..0dff1ef908e2 100644
--- a/python/testData/stubs/StubStructure.py
+++ b/python/testData/stubs/StubStructure.py
@@ -2,7 +2,7 @@ def deco(fun):
   return fun # valid
 
 class FooClass:
-  staticField = 1
+  staticField = deco
   def __init__(self):
     self.instanceField = 2
 
diff --git a/python/testSrc/com/jetbrains/python/PyStubsTest.java b/python/testSrc/com/jetbrains/python/PyStubsTest.java
index 43202c0f65de..d51594ca7889 100644
--- a/python/testSrc/com/jetbrains/python/PyStubsTest.java
+++ b/python/testSrc/com/jetbrains/python/PyStubsTest.java
@@ -45,6 +45,7 @@ public class PyStubsTest extends PyLightFixtureTestCase {
     final PyTargetExpression[] attrs = pyClass.getClassAttributes();
     assertEquals(1, attrs.length);
     assertEquals("staticField", attrs [0].getName());
+    assertTrue(attrs [0].getAssignedQName().matches("deco"));
 
     final PyFunction[] methods = pyClass.getMethods();
     assertEquals(2, methods.length);