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

Closes #713

1 files changed, 84 insertions, 0 deletions

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(
             """