Don't resolve to file-level '__metaclass__' in PyClass.getMetaClassExpression() (PY-12127)

This method is invoked during indexing and it shouldn't resolve
references extensively. Related to this issue, the return type of
PyClass.getMetaClassType() has been changed in order to be able to
return weak metaclass types.

python/psi-api/src/com/jetbrains/python/psi/PyClass.java

@@ -25,6 +25,7 @@ import com.intellij.util.Processor;
 import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
 import com.jetbrains.python.psi.stubs.PyClassStub;
 import com.jetbrains.python.psi.types.PyClassLikeType;
+import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NonNls;
 import org.jetbrains.annotations.NotNull;
@@ -221,14 +222,15 @@ public interface PyClass extends PsiNameIdentifierOwner, PyStatement, NameDefine
   //TODO: Add "addMetaClass" or move methods out of here
   /**
    * Returns the type representing the metaclass of the class if it is explicitly set, null otherwise.
+   *
+   * The metaclass might be defined outside the class in case of Python 2 file-level __metaclass__ attributes.
    */
   @Nullable
-  PyClassLikeType getMetaClassType(@NotNull TypeEvalContext context);
+  PyType getMetaClassType(@NotNull TypeEvalContext context);
 
   /**
    * Returns the expression that defines the metaclass of the class.
    *
-   * It might be defined outside the class in case of file-level __metaclass__ attributes.
    * Operates at the AST level.
    */
   @Nullable

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

@@ -1172,13 +1172,13 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
 
   @Nullable
   @Override
-  public PyClassLikeType getMetaClassType(@NotNull TypeEvalContext context) {
+  public PyType getMetaClassType(@NotNull TypeEvalContext context) {
     if (context.maySwitchToAST(this)) {
       final PyExpression expression = getMetaClassExpression();
       if (expression != null) {
         final PyType type = context.getType(expression);
-        if (type instanceof PyClassLikeType) {
-          return (PyClassLikeType)type;
+        if (type != null) {
+          return type;
         }
       }
     }
@@ -1193,6 +1193,17 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
         }
       }
     }
+    final LanguageLevel level = LanguageLevel.forElement(this);
+    if (level.isOlderThan(LanguageLevel.PYTHON30)) {
+      final PsiFile file = getContainingFile();
+      if (file instanceof PyFile) {
+        final PyFile pyFile = (PyFile)file;
+        final PsiElement element = pyFile.getElementNamed(PyNames.DUNDER_METACLASS);
+        if (element instanceof PyTypedElement) {
+          return context.getType((PyTypedElement)element);
+        }
+      }
+    }
     return null;
   }
 
@@ -1216,14 +1227,6 @@ public class PyClassImpl extends PyPresentableElementImpl<PyClassStub> implement
       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;
   }

python/src/com/jetbrains/python/psi/types/PyClassTypeImpl.java

@@ -267,9 +267,9 @@ public class PyClassTypeImpl extends UserDataHolderBase implements PyClassType {
   @Nullable
   @Override
   public PyClassLikeType getMetaClassType(@NotNull TypeEvalContext context, boolean inherited) {
-    final PyClassLikeType ownMeta = myClass.getMetaClassType(context);
+    final PyType ownMeta = myClass.getMetaClassType(context);
     if (ownMeta != null) {
-      return ownMeta;
+      return (ownMeta instanceof PyClassLikeType) ? (PyClassLikeType)ownMeta : null;
     }
     if (inherited) {
       for (PyClassLikeType ancestor : myClass.getAncestorTypes(context)) {

python/src/com/jetbrains/python/refactoring/classes/membersManager/MethodsManager.java

@@ -9,7 +9,7 @@ import com.jetbrains.python.PyNames;
 import com.jetbrains.python.codeInsight.imports.AddImportHelper;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyFunctionBuilder;
-import com.jetbrains.python.psi.types.PyClassLikeType;
+import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import com.jetbrains.python.refactoring.classes.PyClassRefactoringUtil;
 import org.jetbrains.annotations.NotNull;
@@ -98,7 +98,7 @@ class MethodsManager extends MembersManager<PyFunction> {
   // TODO: Copy/Paste with PyClass.getMeta..
   private static boolean addMetaAbcIfNeeded(@NotNull final PyClass aClass) {
     final PsiFile file = aClass.getContainingFile();
-    final PyClassLikeType type = aClass.getMetaClassType(TypeEvalContext.userInitiated(file));
+    final PyType type = aClass.getMetaClassType(TypeEvalContext.userInitiated(file));
     if (type != null) {
       return false; //User already has metaclass. He probably knows about metaclasses, so we should not add ABCMeta
     }

python/testData/stubs/MetaClass.py

@@ -0,0 +1,13 @@
+class M(type):
+    pass
+
+
+__metaclass__ = M
+
+
+class C(object):
+    __metaclass__ = type
+
+
+class D(object):
+    pass

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

@@ -29,16 +29,17 @@ import com.intellij.psi.impl.source.PsiFileImpl;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.stubs.StubUpdatingIndex;
+import com.intellij.psi.util.QualifiedName;
 import com.intellij.testFramework.TestDataPath;
 import com.intellij.util.indexing.FileBasedIndex;
 import com.jetbrains.python.fixtures.PyTestCase;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyFileImpl;
-import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.psi.impl.PythonLanguageLevelPusher;
 import com.jetbrains.python.psi.stubs.PyClassNameIndex;
 import com.jetbrains.python.psi.stubs.PyClassStub;
 import com.jetbrains.python.psi.stubs.PyVariableNameIndex;
+import com.jetbrains.python.psi.types.TypeEvalContext;
 import com.jetbrains.python.toolbox.Maybe;
 
 import java.util.Collection;
@@ -411,4 +412,16 @@ public class PyStubsTest extends PyTestCase {
     final String docString = foo.getDocStringValue();
     assertEquals("Foo docstring.", docString);
   }
+
+  public void testMetaClass() {
+    final PyFile file = getTestFile();
+    final PyClass c = file.findTopLevelClass("C");
+    assertNotNull(c);
+    assertNotNull(c.getMetaClassExpression());
+    final PyClass d = file.findTopLevelClass("D");
+    assertNotNull(d);
+    assertNull(d.getMetaClassExpression());
+    assertNotNull(d.getMetaClassType(TypeEvalContext.codeInsightFallback()));
+    assertNotParsed(file);
+  }
 }