From 919ba06dd71589c1b074e0e57fccc05534bc765b Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Tue, 1 Nov 2016 15:43:08 +0300
Subject: [PATCH] Support six.add_metaclass decorator

---
 .../python/psi/impl/PyClassImpl.java          | 49 +++++++++++++++++++
 .../python/PythonCompletionTest.java          | 28 +++++++++++
 2 files changed, 77 insertions(+)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java b/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java
index 38df49a07383..8fd6b672f542 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyClassImpl.java
@@ -1455,9 +1455,58 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
       }
     }
 
+    final PyDecoratorList decoratorList = getDecoratorList();
+    if (decoratorList != null) {
+      for (PyDecorator decorator : decoratorList.getDecorators()) {
+        if (isSixAddMetaclass(decorator)) {
+          final PyExpression[] arguments = decorator.getArguments();
+          if (arguments.length != 0) {
+            return arguments[0];
+          }
+        }
+      }
+    }
+
     return null;
   }
 
+  private static boolean isSixAddMetaclass(@NotNull PyDecorator decorator) {
+    final PyExpression callee = decorator.getCallee();
+    if (callee != null && "add_metaclass".equals(callee.getName())) {
+      // SUPPORTED CASES:
+
+      // import six
+      // six.add_metaclass(...)
+
+      // from six import add_metaclass
+      // add_metaclass(...)
+      return true;
+    }
+
+    if (callee instanceof PyReferenceExpression) {
+      // SUPPORTED CASES:
+
+      // from six import add_metaclass as a_m
+      // a_m(...)
+
+      final boolean importedAddMetaclass = StreamEx
+        .of(PyResolveUtil.resolveLocally((PyReferenceExpression)callee))
+        .select(PyImportElement.class)
+        .map(PyImportElement::getImportedQName)
+        .nonNull()
+        .map(QualifiedName::getLastComponent)
+        .nonNull()
+        .findAny("add_metaclass"::equals)
+        .isPresent();
+
+      if (importedAddMetaclass) {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
   @NotNull
   private List<PyClassLikeType> getMROAncestorTypes(@NotNull TypeEvalContext context) throws MROException {
     PyPsiUtils.assertValid(this);
diff --git a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
index 7e8285f71a7e..b820ec4e6a3c 100644
--- a/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
+++ b/python/testSrc/com/jetbrains/python/PythonCompletionTest.java
@@ -1076,6 +1076,34 @@ public class PythonCompletionTest extends PyTestCase {
                          });
   }
 
+  public void testSixAddMetaclass() {
+    final List<String> suggested = doTestByText("import six\n" +
+                                                "class M(type):\n" +
+                                                "    def baz(self):\n" +
+                                                "        pass\n" +
+                                                "@six.add_metaclass(M)\n" +
+                                                "class C(object):\n" +
+                                                "    def foo(self):\n" +
+                                                "        C.ba<caret>()");
+
+    assertNotNull(suggested);
+    assertContainsElements(suggested, "baz");
+  }
+
+  public void testSixAddMetaclassWithAs() {
+    final List<String> suggested = doTestByText("from six import add_metaclass as a_m\n" +
+                                                "class M(type):\n" +
+                                                "    def baz(self):\n" +
+                                                "        pass\n" +
+                                                "@a_m(M)\n" +
+                                                "class C(object):\n" +
+                                                "    def foo(self):\n" +
+                                                "        C.ba<caret>()");
+
+    assertNotNull(suggested);
+    assertContainsElements(suggested, "baz");
+  }
+
   @Override
   protected String getTestDataPath() {
     return super.getTestDataPath() + "/completion";