From a488feb93116cea555302761014a3309d76d8793 Mon Sep 17 00:00:00 2001
From: Claudiu Popa <pcmanticore@gmail.com>
Date: Sat, 31 Oct 2015 21:31:36 +0200
Subject: Fix some inconsistencies with accessing the first argument of methods

If the current class is a metaclass (inherits from `type`), then
the first argument is always a class, not an instance.
---
 astroid/arguments.py                | 10 +++++++++
 astroid/protocols.py                | 10 ++++++---
 astroid/tests/unittest_inference.py | 42 +++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/astroid/arguments.py b/astroid/arguments.py
index 5670fa8..e80f5f4 100644
--- a/astroid/arguments.py
+++ b/astroid/arguments.py
@@ -183,6 +183,16 @@ class CallSite(object):
                 else:
                     # XXX can do better ?
                     boundnode = funcnode.parent.frame()
+                
+                if isinstance(boundnode, nodes.ClassDef):
+                    # Verify that we're accessing a method
+                    # of the metaclass through a class, as in
+                    # `cls.metaclass_method`. In this case, the
+                    # first argument is always the class. 
+                    method_scope = funcnode.parent.scope()
+                    if method_scope is boundnode.metaclass():
+                        return iter((boundnode, ))
+
                 if funcnode.type == 'method':
                     if not isinstance(boundnode, bases.Instance):
                         boundnode = bases.Instance(boundnode)
diff --git a/astroid/protocols.py b/astroid/protocols.py
index 7c9e4b9..28b1dd2 100644
--- a/astroid/protocols.py
+++ b/astroid/protocols.py
@@ -280,12 +280,16 @@ def _arguments_infer_argname(self, name, context):
     # first argument of instance/class method
     if self.args and getattr(self.args[0], 'name', None) == name:
         functype = self.parent.type
+        cls = self.parent.parent.scope()
+        is_metaclass = isinstance(cls, nodes.ClassDef) and cls.type == 'metaclass'
+        # If this is a metaclass, then the first argument will always
+        # be the class, not an instance.
+        if is_metaclass or functype == 'classmethod':
+            yield cls
+            return
         if functype == 'method':
             yield bases.Instance(self.parent.parent.frame())
             return
-        if functype == 'classmethod':
-            yield self.parent.parent.frame()
-            return
 
     if context and context.callcontext:
         call_site = arguments.CallSite(context.callcontext)
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):
 
-- 
cgit v1.2.1