Add support of 'six.with_metaclass' (#841)

Closes #713

3 files changed, 127 insertions, 0 deletions

diff --git a/ChangeLog b/ChangeLog
index 9ff42a6e..fd0c428a 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -94,6 +94,10 @@ Release Date: TBA
 
   Fixes #843
 
+* Fix interpretation of ``six.with_metaclass`` class definitions.
+
+  Fixes #713
+
 * Reduce memory usage of astroid's module cache.
 
 * Remove dependency on `imp`.
diff --git a/astroid/brain/brain_six.py b/astroid/brain/brain_six.py
index 389037f2..a998213f 100644
--- a/astroid/brain/brain_six.py
+++ b/astroid/brain/brain_six.py
@@ -22,6 +22,7 @@ from astroid import nodes
 
 
 SIX_ADD_METACLASS = "six.add_metaclass"
+SIX_WITH_METACLASS = "six.with_metaclass"
 
 
 def _indent(text, prefix, predicate=None):
@@ -190,6 +191,39 @@ def transform_six_add_metaclass(node):
             return node
 
 
+def _looks_like_nested_from_six_with_metaclass(node):
+    if len(node.bases) != 1:
+        return False
+    base = node.bases[0]
+    if not isinstance(base, nodes.Call):
+        return False
+    try:
+        if hasattr(base.func, "expr"):
+            # format when explicit 'six.with_metaclass' is used
+            mod = base.func.expr.name
+            func = base.func.attrname
+            func = "{}.{}".format(mod, func)
+        else:
+            # format when 'with_metaclass' is used directly (local import from six)
+            # check reference module to avoid 'with_metaclass' name clashes
+            mod = base.parent.parent
+            import_from = mod.locals["with_metaclass"][0]
+            func = "{}.{}".format(import_from.modname, base.func.name)
+    except (AttributeError, KeyError, IndexError):
+        return False
+    return func == SIX_WITH_METACLASS
+
+
+def transform_six_with_metaclass(node):
+    """Check if the given class node is defined with *six.with_metaclass*
+
+    If so, inject its argument as the metaclass of the underlying class.
+    """
+    call = node.bases[0]
+    node._metaclass = call.args[0]
+    node.bases = call.args[1:]
+
+
 register_module_extender(MANAGER, "six", six_moves_transform)
 register_module_extender(
     MANAGER, "requests.packages.urllib3.packages.six", six_moves_transform
@@ -200,3 +234,8 @@ MANAGER.register_transform(
     transform_six_add_metaclass,
     _looks_like_decorated_with_six_add_metaclass,
 )
+MANAGER.register_transform(
+    nodes.ClassDef,
+    transform_six_with_metaclass,
+    _looks_like_nested_from_six_with_metaclass,
+)
diff --git a/tests/unittest_inference.py b/tests/unittest_inference.py
index 7b80b530..1512456e 100644
--- a/tests/unittest_inference.py
+++ b/tests/unittest_inference.py
@@ -2997,6 +2997,23 @@ class InferenceTest(resources.SysPathSetup, unittest.TestCase):
         self.assertIsInstance(inferred, nodes.Const)
         self.assertEqual(inferred.value, 24)
 
+    def test_with_metaclass__getitem__(self):
+        ast_node = extract_node(
+            """
+        class Meta(type):
+            def __getitem__(cls, arg):
+                return 24
+        import six
+        class A(six.with_metaclass(Meta)):
+            pass
+
+        A['Awesome'] #@
+        """
+        )
+        inferred = next(ast_node.infer())
+        self.assertIsInstance(inferred, nodes.Const)
+        self.assertEqual(inferred.value, 24)
+
     def test_bin_op_classes(self):
         ast_node = extract_node(
             """
@@ -3015,6 +3032,23 @@ class InferenceTest(resources.SysPathSetup, unittest.TestCase):
         self.assertIsInstance(inferred, nodes.Const)
         self.assertEqual(inferred.value, 24)
 
+    def test_bin_op_classes_with_metaclass(self):
+        ast_node = extract_node(
+            """
+        class Meta(type):
+            def __or__(self, other):
+                return 24
+        import six
+        class A(six.with_metaclass(Meta)):
+            pass
+
+        A | A
+        """
+        )
+        inferred = next(ast_node.infer())
+        self.assertIsInstance(inferred, nodes.Const)
+        self.assertEqual(inferred.value, 24)
+
     def test_bin_op_supertype_more_complicated_example(self):
         ast_node = extract_node(
             """
@@ -3354,6 +3388,22 @@ class InferenceTest(resources.SysPathSetup, unittest.TestCase):
         self.assertIsInstance(inferred, nodes.Const)
         self.assertEqual(inferred.value, 42)
 
+    def test_unary_op_classes_with_metaclass(self):
+        ast_node = extract_node(
+            """
+        import six
+        class Meta(type):
+            def __invert__(self):
+                return 42
+        class A(six.with_metaclass(Meta)):
+            pass
+        ~A
+        """
+        )
+        inferred = next(ast_node.infer())
+        self.assertIsInstance(inferred, nodes.Const)
+        self.assertEqual(inferred.value, 42)
+
     def _slicing_test_helper(self, pairs, cls, get_elts):
         for code, expected in pairs:
             ast_node = extract_node(code)
@@ -3750,6 +3800,40 @@ class InferenceTest(resources.SysPathSetup, unittest.TestCase):
         self.assertIsInstance(inferred, nodes.ClassDef)
         self.assertEqual(inferred.name, "B")
 
+    def test_With_metaclass_subclasses_arguments_are_classes_not_instances(self):
+        ast_node = extract_node(
+            """
+        class A(type):
+            def test(cls):
+                return cls
+        import six
+        class B(six.with_metaclass(A)):
+            pass
+
+        B.test() #@
+        """
+        )
+        inferred = next(ast_node.infer())
+        self.assertIsInstance(inferred, nodes.ClassDef)
+        self.assertEqual(inferred.name, "B")
+
+    def test_With_metaclass_with_partial_imported_name(self):
+        ast_node = extract_node(
+            """
+        class A(type):
+            def test(cls):
+                return cls
+        from six import with_metaclass
+        class B(with_metaclass(A)):
+            pass
+
+        B.test() #@
+        """
+        )
+        inferred = next(ast_node.infer())
+        self.assertIsInstance(inferred, nodes.ClassDef)
+        self.assertEqual(inferred.name, "B")
+
     def test_infer_cls_in_class_methods(self):
         ast_nodes = extract_node(
             """