Merge bookmark heads.

1 files changed, 42 insertions, 0 deletions

diff --git a/astroid/tests/unittest_inference.py b/astroid/tests/unittest_inference.py
index fa090f8..0f45d3a 100644
--- a/astroid/tests/unittest_inference.py
+++ b/astroid/tests/unittest_inference.py
@@ -3143,6 +3143,48 @@ class InferenceTest(resources.SysPathSetup, unittest.TestCase):
         self.assertIsInstance(third_c, Instance)
         self.assertEqual(third_c.name, 'A')
 
+    def test_metaclass_subclasses_arguments_are_classes_not_instances(self):
+        ast_node = test_utils.extract_node('''
+        class A(type):
+            def test(cls):
+                return cls        
+        import six
+        @six.add_metaclass(A)
+        class B(object):
+            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 = test_utils.extract_node('''
+        class A(type):
+            def __call__(cls):
+                cls #@
+        class B(object):
+            def __call__(cls):
+                cls #@        
+        ''')
+        first = next(ast_nodes[0].infer())
+        self.assertIsInstance(first, nodes.ClassDef)
+        second = next(ast_nodes[1].infer())
+        self.assertIsInstance(second, Instance)
+
+    @unittest.expectedFailure
+    def test_metaclass_arguments_are_classes_not_instances(self):
+        ast_node = test_utils.extract_node('''
+        class A(type):
+            def test(cls): return cls
+        A.test() #@
+        ''')
+        # This is not supported yet
+        inferred = next(ast_node.infer())
+        self.assertIsInstance(inferred, ClassDef)
+        self.assertEqual(inferred.name, 'A')
+
 
 class GetattrTest(unittest.TestCase):