to avoid unnecessary tree loading in Django code, when the value assigned to a target expression is a call, store the callee qualified name in stubs

python/src/com/jetbrains/python/psi/PyFileElementType.java

@@ -41,7 +41,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
 
   @Override
   public int getStubVersion() {
-    return 27;
+    return 28;
   }
 
   @Override

python/src/com/jetbrains/python/psi/PyTargetExpression.java

@@ -27,6 +27,15 @@ public interface PyTargetExpression extends PyQualifiedExpression, PsiNamedEleme
   @Nullable
   PyQualifiedName getAssignedQName();
 
+  /**
+   * If the value assigned to the target expression is a call, returns the (unqualified and unresolved) name of the
+   * callee. Otherwise, returns null.
+   *
+   * @return the name of the callee or null if the assigned value is not a call.
+   */
+  @Nullable
+  PyQualifiedName getCalleeName();
+
   @NotNull
   PsiReference getReference();
 }

python/src/com/jetbrains/python/psi/PyUtil.java

@@ -669,10 +669,10 @@ public class PyUtil {
   }
 
   @Nullable
-  public static PyExpression getKeywordArgument(PyCallExpressionImpl expr, String keyword) {
+  public static PyExpression getKeywordArgument(PyCallExpression expr, String keyword) {
     for (PyExpression arg : expr.getArguments()) {
-      if (arg instanceof PyKeywordArgumentImpl) {
-        PyKeywordArgumentImpl kwarg = (PyKeywordArgumentImpl)arg;
+      if (arg instanceof PyKeywordArgument) {
+        PyKeywordArgument kwarg = (PyKeywordArgument)arg;
         if (keyword.equals(kwarg.getKeyword())) {
           return kwarg.getValueExpression();
         }

python/src/com/jetbrains/python/psi/impl/PyQualifiedName.java

@@ -85,6 +85,10 @@ public class PyQualifiedName {
     return true;
   }
 
+  public boolean endsWith(@NotNull String suffix) {
+    return suffix.equals(getLastComponent());
+  }
+
   public static void serialize(@Nullable PyQualifiedName qName, StubOutputStream dataStream) throws IOException {
     if (qName == null) {
       dataStream.writeVarInt(0);

python/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java

@@ -214,12 +214,32 @@ public class PyTargetExpressionImpl extends PyPresentableElementImpl<PyTargetExp
   public PyQualifiedName getAssignedQName() {
     final PyTargetExpressionStub stub = getStub();
     if (stub != null) {
-      return stub.getInitializer();
+      if (stub.getInitializerType() == PyTargetExpressionStub.InitializerType.ReferenceExpression) {
+        return stub.getInitializer();
+      }
+      return null;
     }
     final PyExpression value = findAssignedValue();
     return value instanceof PyReferenceExpression ? ((PyReferenceExpression) value).asQualifiedName() : null;
   }
 
+  @Override
+  public PyQualifiedName getCalleeName() {
+    final PyTargetExpressionStub stub = getStub();
+    if (stub != null) {
+      if (stub.getInitializerType() == PyTargetExpressionStub.InitializerType.CallExpression) {
+        return stub.getInitializer();
+      }
+      return null;
+    }
+    final PyExpression value = findAssignedValue();
+    if (value instanceof PyCallExpression) {
+      final PyExpression callee = ((PyCallExpression)value).getCallee();
+      return callee instanceof PyReferenceExpression ? ((PyReferenceExpression) callee).asQualifiedName() : null;
+    }
+    return null;
+  }
+
   @NotNull
   @Override
   public PsiReference getReference() {

python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java

@@ -40,26 +40,33 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
       return new PyTargetExpressionStubImpl(name, prop, parentStub);
     }
     else {
-      final PyQualifiedName initializer = assignedValue instanceof PyReferenceExpression
-                                          ? ((PyReferenceExpression) assignedValue).asQualifiedName()
-                                          : null;
-      return new PyTargetExpressionStubImpl(name, initializer, parentStub);
+      PyTargetExpressionStub.InitializerType initializerType = PyTargetExpressionStub.InitializerType.Other;
+      PyQualifiedName initializer = null;
+      if (assignedValue instanceof PyReferenceExpression) {
+        initializerType = PyTargetExpressionStub.InitializerType.ReferenceExpression;
+        initializer = ((PyReferenceExpression) assignedValue).asQualifiedName();
+      }
+      else if (assignedValue instanceof PyCallExpression) {
+        initializerType = PyTargetExpressionStub.InitializerType.CallExpression;
+        final PyExpression callee = ((PyCallExpression)assignedValue).getCallee();
+        if (callee instanceof PyReferenceExpression) {
+          initializer = ((PyReferenceExpression) callee).asQualifiedName();
+        }
+      }
+      return new PyTargetExpressionStubImpl(name, initializerType, initializer, parentStub);
     }
   }
 
-  private static final int SIMPLE = 0; // stream stores a PyTargetExpressionStubImpl
-  private static final int PROPERTY = 1; // stream stores a PyTargetExpressionPropertyStubImpl
-
   public void serialize(final PyTargetExpressionStub stub, final StubOutputStream stream)
       throws IOException {
     stream.writeName(stub.getName());
-    PropertyStubStorage prop = stub.getPropertyPack();
-    if (prop != null) {
-      stream.writeVarInt(PROPERTY);
-      prop.serialize(stream);
+    stream.writeVarInt(stub.getInitializerType().getIndex());
+    if (stub.getInitializerType() == PyTargetExpressionStub.InitializerType.Property) {
+      final PropertyStubStorage propertyPack = stub.getPropertyPack();
+      assert propertyPack != null;
+      propertyPack.serialize(stream);
     }
     else {
-      stream.writeVarInt(SIMPLE);
       PyQualifiedName.serialize(stub.getInitializer(), stream);
     }
   }
@@ -67,19 +74,13 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
   public PyTargetExpressionStub deserialize(final StubInputStream stream, final StubElement parentStub)
       throws IOException {
     String name = StringRef.toString(stream.readName());
-    int code = stream.readVarInt();
-    if (code == SIMPLE) {
-      PyQualifiedName initializer = PyQualifiedName.deserialize(stream);
-      return new PyTargetExpressionStubImpl(name, initializer, parentStub);
-    }
-    else if (code == PROPERTY) {
+    PyTargetExpressionStub.InitializerType initializerType = PyTargetExpressionStub.InitializerType.fromIndex(stream.readVarInt());
+    if (initializerType == PyTargetExpressionStub.InitializerType.Property) {
       PropertyStubStorage prop = PropertyStubStorage.deserialize(stream);
       return new PyTargetExpressionStubImpl(name, prop, parentStub);
     }
-    else {
-      assert false : "Unknown code in stream: " + code;
-      return null; // to keep inspections safe
-    }
+    PyQualifiedName initializer = PyQualifiedName.deserialize(stream);
+    return new PyTargetExpressionStubImpl(name, initializerType, initializer, parentStub);
   }
 
   public boolean shouldCreateStub(final ASTNode node) {

python/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionStubImpl.java

@@ -14,6 +14,7 @@ import org.jetbrains.annotations.Nullable;
  */
 public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> implements PyTargetExpressionStub {
   private final String myName;
+  private final InitializerType myInitializerType;
   private final PyQualifiedName myInitializer;
 
   private final PropertyStubStorage myPropertyPack;
@@ -21,13 +22,17 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
   public PyTargetExpressionStubImpl(String name, PropertyStubStorage propertyPack, StubElement parent) {
     super(parent, PyElementTypes.TARGET_EXPRESSION);
     myName = name;
+    myInitializerType = InitializerType.Property;
     myInitializer = null;
     myPropertyPack = propertyPack;
   }
   
-  public PyTargetExpressionStubImpl(final String name, final PyQualifiedName initializer, final StubElement parentStub) {
+  public PyTargetExpressionStubImpl(final String name, final InitializerType initializerType,
+                                    final PyQualifiedName initializer, final StubElement parentStub) {
     super(parentStub, PyElementTypes.TARGET_EXPRESSION);
     myName = name;
+    assert initializerType != InitializerType.Property;
+    myInitializerType = initializerType;
     myInitializer = initializer;
     myPropertyPack = null;
   }
@@ -36,6 +41,10 @@ public class PyTargetExpressionStubImpl extends StubBase<PyTargetExpression> imp
     return myName;
   }
 
+  public InitializerType getInitializerType() {
+    return myInitializerType;
+  }
+
   public PyQualifiedName getInitializer() {
     return myInitializer;
   }

python/src/com/jetbrains/python/psi/stubs/PyTargetExpressionStub.java

@@ -9,6 +9,34 @@ import org.jetbrains.annotations.Nullable;
  * @author yole
  */
 public interface PyTargetExpressionStub extends NamedStub<PyTargetExpression> {
+  enum InitializerType {
+    ReferenceExpression(1),
+    CallExpression(2),
+    Property(3),
+    Other(0);
+
+    private int myIndex;
+
+    InitializerType(int index) {
+      myIndex = index;
+    }
+
+    public int getIndex() {
+      return myIndex;
+    }
+
+    public static InitializerType fromIndex(int index) {
+      switch (index) {
+        case 1: return ReferenceExpression;
+        case 2: return CallExpression;
+        case 3: return Property;
+        default: return Other;
+      }
+    }
+  }
+
+  InitializerType getInitializerType();
+
   @Nullable
   PyQualifiedName getInitializer();
 

python/testData/stubs/StubStructure.py

@@ -3,6 +3,7 @@ def deco(fun):
 
 class FooClass:
   staticField = deco
+  globs = globals()
   def __init__(self):
     self.instanceField = 2
 

python/testSrc/com/jetbrains/python/PyStubsTest.java

@@ -50,7 +50,7 @@ public class PyStubsTest extends PyLightFixtureTestCase {
       assertEquals("StubStructure.FooClass", pyClass.getQualifiedName());
 
       final List<PyTargetExpression> attrs = pyClass.getClassAttributes();
-      assertEquals(1, attrs.size());
+      assertEquals(2, attrs.size());
       assertEquals("staticField", attrs.get(0).getName());
       assertTrue(attrs.get(0).getAssignedQName().matches("deco"));