From 1ee1c255cbb2ea106b0a01260488843f86eaed2a Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@gmail.com>
Date: Tue, 21 Jan 2014 00:29:20 +0400
Subject: [PATCH] Extracted getMetaClassExpression()

---
 .../python/psi/impl/PyClassImpl.java          | 68 +++++++++----------
 1 file changed, 31 insertions(+), 37 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java b/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java
index 16623a233a5c..7ebff88ca682 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java
@@ -1170,49 +1170,43 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
   @Nullable
   @Override
   public PyClassLikeType getMetaClassType(@NotNull TypeEvalContext context) {
-    final LanguageLevel level = LanguageLevel.forElement(this);
-    if (level.isAtLeast(LanguageLevel.PYTHON30)) {
-      // TODO: This requires switching from stubs to AST
-      final PyExpression[] superClassExpressions = getSuperClassExpressions();
-      for (PyExpression superClassExpression : superClassExpressions) {
-        if (superClassExpression instanceof PyKeywordArgument) {
-          final PyKeywordArgument argument = (PyKeywordArgument)superClassExpression;
-          if (PyNames.METACLASS.equals(argument.getKeyword())) {
-            final PyClassLikeType type = getClassLikeType(argument, context);
-            if (type != null) {
-              return type;
-            }
-          }
-        }
-      }
-    }
-    else {
-      final PyTargetExpression metaClassAttribute = findClassAttribute(PyNames.DUNDER_METACLASS, false);
-      if (metaClassAttribute != null) {
-        final PyClassLikeType type = getClassLikeType(metaClassAttribute, context);
-        if (type != null) {
-          return type;
-        }
-      }
-      final PsiFile containingFile = getContainingFile();
-      if (containingFile instanceof PyFile) {
-        final PsiElement element = ((PyFile)containingFile).getElementNamed(PyNames.DUNDER_METACLASS);
-        if (element instanceof PyTypedElement) {
-          final PyClassLikeType type = getClassLikeType((PyTypedElement)element, context);
-          if (type != null) {
-            return type;
-          }
-        }
+    final PyExpression expression = getMetaClassExpression();
+    if (expression != null) {
+      final PyType type = context.getType(expression);
+      if (type instanceof PyClassLikeType) {
+        return (PyClassLikeType)type;
       }
     }
     return null;
   }
 
   @Nullable
-  private static PyClassLikeType getClassLikeType(@NotNull PyTypedElement element, @NotNull TypeEvalContext context) {
-    final PyType type = context.getType(element);
-    if (type instanceof PyClassLikeType) {
-      return (PyClassLikeType)type;
+  public PyExpression getMetaClassExpression() {
+    final LanguageLevel level = LanguageLevel.forElement(this);
+    if (level.isAtLeast(LanguageLevel.PYTHON30)) {
+      // Requires AST access
+      for (PyExpression expression : getSuperClassExpressions()) {
+        if (expression instanceof PyKeywordArgument) {
+          final PyKeywordArgument argument = (PyKeywordArgument)expression;
+          if (PyNames.METACLASS.equals(argument.getKeyword())) {
+            return argument.getValueExpression();
+          }
+        }
+      }
+    }
+    else {
+      final PyTargetExpression attribute = findClassAttribute(PyNames.DUNDER_METACLASS, false);
+      if (attribute != null) {
+        return attribute;
+      }
+      final PsiFile file = getContainingFile();
+      if (file instanceof PyFile) {
+        final PyFile pyFile = (PyFile)file;
+        final PsiElement element = pyFile.getElementNamed(PyNames.DUNDER_METACLASS);
+        if (element instanceof PyExpression) {
+          return (PyExpression)element;
+        }
+      }
     }
     return null;
   }