add support to 'six.with_metaclass' (fix pylint-dev#713)

astroid/brain/brain_six.py

@@ -189,6 +189,29 @@ 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:
+        func = next(base.func.infer())
+    except InferenceError:
+        return False
+    return func.qname() == '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
@@ -199,3 +222,8 @@ def transform_six_add_metaclass(node):
     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,
+)

tests/unittest_inference.py

@@ -3001,6 +3001,23 @@ class A(object):
         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(
             """
@@ -3019,6 +3036,23 @@ class A(object):
         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(
             """
@@ -3358,6 +3392,22 @@ class A(object):
         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)
@@ -3753,6 +3803,23 @@ class B(object):
         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_infer_cls_in_class_methods(self):
         ast_nodes = extract_node(
             """