From 0eb783c7c845371db30690a871f37ae4a49b776c Mon Sep 17 00:00:00 2001
From: Dmitry Jemerov <yole@jetbrains.com>
Date: Fri, 2 Mar 2012 19:49:46 +0100
Subject: [PATCH] support __metaclass__ attribute on new-style classes
 (PY-1392)

---
 .../python/psi/types/PyClassType.java          | 18 ++++++++++++++++--
 python/testData/resolve/Metaclass.py           | 10 ++++++++++
 .../com/jetbrains/python/PyResolveTest.java    |  5 +++++
 3 files changed, 31 insertions(+), 2 deletions(-)
 create mode 100644 python/testData/resolve/Metaclass.py

diff --git a/python/src/com/jetbrains/python/psi/types/PyClassType.java b/python/src/com/jetbrains/python/psi/types/PyClassType.java
index d69580233a08..1435b0ee8c9c 100644
--- a/python/src/com/jetbrains/python/psi/types/PyClassType.java
+++ b/python/src/com/jetbrains/python/psi/types/PyClassType.java
@@ -177,7 +177,7 @@ public class PyClassType extends UserDataHolderBase implements PyCallableType {
       }
     }
     if (isDefinition() && myClass.isNewStyleClass()) {
-      PyClassType typeType = PyBuiltinCache.getInstance(myClass).getObjectType("type");
+      PyClassType typeType = getMetaclassType();
       if (typeType != null) {
         List<? extends RatedResolveResult> typeMembers = typeType.resolveMember(name, location, direction, resolveContext);
         if (typeMembers != null && !typeMembers.isEmpty()) {
@@ -188,6 +188,20 @@ public class PyClassType extends UserDataHolderBase implements PyCallableType {
     return Collections.emptyList();
   }
 
+  private PyClassType getMetaclassType() {
+    final PyTargetExpression metaClassAttribute = myClass.findClassAttribute(PyNames.METACLASS, true);
+    if (metaClassAttribute != null) {
+      final PyExpression metaclass = metaClassAttribute.findAssignedValue();
+      if (metaclass instanceof PyReferenceExpression) {
+        final QualifiedResolveResult result = ((PyReferenceExpression)metaclass).followAssignmentsChain(PyResolveContext.noImplicits());
+        if (result.getElement() instanceof PyClass) {
+          return new PyClassType((PyClass) result.getElement(), false);
+        }
+      }
+    }
+    return PyBuiltinCache.getInstance(myClass).getObjectType("type");
+  }
+
   @Override
   public PyType getCallType() {
     if (isDefinition()) {
@@ -265,7 +279,7 @@ public class PyClassType extends UserDataHolderBase implements PyCallableType {
     }
 
     if (isDefinition() && myClass.isNewStyleClass()) {
-      PyClassType typeType = PyBuiltinCache.getInstance(myClass).getObjectType("type");
+      PyClassType typeType = getMetaclassType();
       if (typeType != null) {
         Collections.addAll(ret, typeType.getCompletionVariants(prefix, location, context));
       }
diff --git a/python/testData/resolve/Metaclass.py b/python/testData/resolve/Metaclass.py
new file mode 100644
index 000000000000..53898502a0ba
--- /dev/null
+++ b/python/testData/resolve/Metaclass.py
@@ -0,0 +1,10 @@
+class PluginMetaclass(type):
+    def getStore(cls):
+        pass
+
+class Plugin(object):
+    __metaclass__ = PluginMetaclass
+
+    def foo(self):
+        Plugin.getStore()
+#                    <ref>
\ No newline at end of file
diff --git a/python/testSrc/com/jetbrains/python/PyResolveTest.java b/python/testSrc/com/jetbrains/python/PyResolveTest.java
index 1d986cb42e2c..edb189bfdc51 100644
--- a/python/testSrc/com/jetbrains/python/PyResolveTest.java
+++ b/python/testSrc/com/jetbrains/python/PyResolveTest.java
@@ -445,4 +445,9 @@ public class PyResolveTest extends PyResolveTestCase {
     PyExpression value = target.findAssignedValue();
     assertInstanceOf(value, PyNumericLiteralExpression.class);
   }
+
+  public void testMetaclass() {
+    final PyFunction function = assertResolvesTo(PyFunction.class, "getStore");
+    assertEquals("PluginMetaclass", function.getContainingClass().getName());
+  }
 }