Merged in PCManticore/astroid/metaclass (pull request #9)

Metaclass support

6 files changed, 171 insertions, 6 deletions

diff --git a/ChangeLog b/ChangeLog
index 4df1b16..d487979 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -2,6 +2,9 @@ Change log for the astroid package (used to be astng)
 =====================================================
 
 --
+    * Add a `metaclass` function to `Class` nodes to 
+      retrieve their metaclass.
+ 
     * Add support for inferring arguments to namedtuple invocations.
 
     * Make sure that objects returned for namedtuple
diff --git a/as_string.py b/as_string.py
index fcff19e..b7e6257 100644
--- a/as_string.py
+++ b/as_string.py
@@ -144,7 +144,17 @@ class AsStringVisitor(object):
         """return an astroid.Class node as string"""
         decorate = node.decorators and node.decorators.accept(self)  or ''
         bases =  ', '.join([n.accept(self) for n in node.bases])
-        bases = bases and '(%s)' % bases or ''
+        if sys.version_info[0] == 2:
+            bases = bases and '(%s)' % bases or ''
+        else:
+            metaclass = node.metaclass()
+            if metaclass:
+                if bases:
+                    bases = '(%s, metaclass=%s)' % (bases, metaclass.name)
+                else:
+                    bases = '(metaclass=%s)' % metaclass.name
+            else:
+                bases = bases and '(%s)' % bases or ''
         docs = node.doc and '\n%s"""%s"""' % (INDENT, node.doc) or ''
         return '\n\n%sclass %s%s:%s\n%s\n' % (decorate, node.name, bases, docs,
                                             self._stmt_list( node.body))
diff --git a/rebuilder.py b/rebuilder.py
index 7f4c7a7..c9d2e21 100644
--- a/rebuilder.py
+++ b/rebuilder.py
@@ -21,7 +21,7 @@ order to get a single Astroid representation
 
 import sys
 from warnings import warn
-from _ast import (Expr as Discard, Str,
+from _ast import (Expr as Discard, Str, Name, Attribute,
     # binary operators
     Add, Div, FloorDiv,  Mod, Mult, Pow, Sub, BitAnd, BitOr, BitXor,
     LShift, RShift,
@@ -114,7 +114,11 @@ def _set_infos(oldnode, newnode, parent):
         newnode.col_offset = oldnode.col_offset
     newnode.set_line_info(newnode.last_child()) # set_line_info accepts None
 
-
+def _infer_metaclass(node):    
+    if isinstance(node, Name):
+        return node.id
+    elif isinstance(node, Attribute):
+        return node.attr
 
 
 class TreeRebuilder(object):
@@ -245,7 +249,7 @@ class TreeRebuilder(object):
                     continue
         elif getattr(newnode.targets[0], 'name', None) == '__metaclass__':
             # XXX check more...
-            self._metaclass[-1] = 'type' # XXX get the actual metaclass
+            self._metaclass[-1] = _infer_metaclass(node.value)
         newnode.set_line_info(newnode.last_child())
         return newnode
 
@@ -331,7 +335,7 @@ class TreeRebuilder(object):
         if not newnode.bases:
             # no base classes, detect new / style old style according to
             # current scope
-            newnode._newstyle = metaclass == 'type'
+            newnode._newstyle = metaclass in ('type', 'ABCMeta')
         newnode.parent.frame().set_local(newnode.name, newnode)
         return newnode
 
@@ -934,6 +938,14 @@ class TreeRebuilder3k(TreeRebuilder):
     def visit_yieldfrom(self, node, parent):
         return self.visit_yield(node, parent)
 
+    def visit_class(self, node, parent):
+        newnode = super(TreeRebuilder3k, self).visit_class(node, parent)
+        for keyword in node.keywords:
+            if keyword.arg == 'metaclass':
+                newnode._metaclass = self.visit(keyword, newnode).value
+                break
+        return newnode
+
 if sys.version_info >= (3, 0):
     TreeRebuilder = TreeRebuilder3k
 
diff --git a/scoped_nodes.py b/scoped_nodes.py
index 0ee29be..a7f6ee8 100644
--- a/scoped_nodes.py
+++ b/scoped_nodes.py
@@ -990,3 +990,25 @@ class Class(Statement, LocalsDictNodeNG, FilterStmtsMixin):
                 yield iface
         if missing:
             raise InferenceError()
+
+    _metaclass = None
+    def metaclass(self):
+        """ Return the metaclass of this class """
+        if self._metaclass:
+            # Expects this from Py3k TreeRebuilder
+            try:
+                return next(self._metaclass.infer())
+            except InferenceError:
+                return 
+
+        try:
+            meta = self.getattr('__metaclass__')[0]
+        except NotFoundError:
+            return
+        try:
+            infered = meta.infer().next()
+        except InferenceError:
+            return
+        if infered is YES: # don't expose this
+            return None
+        return infered
diff --git a/test/unittest_python3.py b/test/unittest_python3.py
index 84650fc..d17b143 100644
--- a/test/unittest_python3.py
+++ b/test/unittest_python3.py
@@ -16,12 +16,14 @@
 # You should have received a copy of the GNU Lesser General Public License along
 # with astroid. If not, see <http://www.gnu.org/licenses/>.
 import sys
+from textwrap import dedent
 
 from logilab.common.testlib import TestCase, unittest_main, require_version
 
 from astroid.node_classes import Assign
 from astroid.manager import AstroidManager
 from astroid.builder import AstroidBuilder
+from astroid.scoped_nodes import Class
 
 
 class Python3TC(TestCase):
@@ -39,5 +41,56 @@ class Python3TC(TestCase):
 
         self.assertTrue(isinstance(node.ass_type(), Assign))
 
+    # metaclass tests
+
+    @require_version('3.0')
+    def test_simple_metaclass(self):
+        astroid = self.builder.string_build("class Test(metaclass=type): pass")
+        klass = astroid.body[0]
+
+        metaclass = klass.metaclass()
+        self.assertIsInstance(metaclass, Class)
+        self.assertEqual(metaclass.name, 'type')
+
+    @require_version('3.0')
+    def test_metaclass_error(self):
+        astroid = self.builder.string_build("class Test(metaclass=typ): pass")
+        klass = astroid.body[0]
+        self.assertFalse(klass.metaclass())
+
+    @require_version('3.0')
+    def test_metaclass_imported(self): 
+        astroid = self.builder.string_build(dedent("""
+        from abc import ABCMeta 
+        class Test(metaclass=ABCMeta): pass"""))
+        klass = astroid.body[1]
+
+        metaclass = klass.metaclass()
+        self.assertIsInstance(metaclass, Class)
+        self.assertEqual(metaclass.name, 'ABCMeta')       
+
+    @require_version('3.0')
+    def test_as_string(self):
+        body = dedent("""
+        from abc import ABCMeta 
+        class Test(metaclass=ABCMeta): pass""")
+        astroid = self.builder.string_build(body)
+        klass = astroid.body[1]
+
+        self.assertEqual(klass.as_string(), 
+                         '\n\nclass Test(metaclass=ABCMeta):\n    pass\n')
+
+    @require_version('3.0')
+    def test_old_syntax_works(self):
+        astroid = self.builder.string_build(dedent("""
+        class Test:
+            __metaclass__ = type
+        class SubTest(Test): pass
+        """))
+        klass = astroid['SubTest']
+        metaclass = klass.metaclass()
+        self.assertIsInstance(metaclass, Class)
+        self.assertEqual(metaclass.name, 'type')
+                         
 if __name__ == '__main__':
     unittest_main()
diff --git a/test/unittest_scoped_nodes.py b/test/unittest_scoped_nodes.py
index a967473..da96b2a 100644
--- a/test/unittest_scoped_nodes.py
+++ b/test/unittest_scoped_nodes.py
@@ -21,8 +21,9 @@ function)
 
 import sys
 from os.path import join, abspath, dirname
+from textwrap import dedent
 
-from logilab.common.testlib import TestCase, unittest_main
+from logilab.common.testlib import TestCase, unittest_main, require_version
 
 from astroid import builder, nodes, scoped_nodes, \
      InferenceError, NotFoundError, NoDefault
@@ -662,6 +663,70 @@ def g2():
         self.assertEqual(astroid['g2'].fromlineno, 9)
         self.assertEqual(astroid['g2'].tolineno, 10)
 
+    def test_simple_metaclass(self):
+        astroid = abuilder.string_build(dedent("""
+        class Test(object):
+            __metaclass__ = type
+        """))
+        klass = astroid['Test']
+
+        metaclass = klass.metaclass()
+        self.assertIsInstance(metaclass, scoped_nodes.Class)
+        self.assertEqual(metaclass.name, 'type')
+
+    def test_metaclass_error(self):
+        astroid = abuilder.string_build(dedent("""
+        class Test(object):
+            __metaclass__ = typ
+        """))
+        klass = astroid['Test']
+        self.assertFalse(klass.metaclass())
+
+    @require_version('2.7')
+    def test_metaclass_imported(self): 
+        astroid = abuilder.string_build(dedent("""
+        from abc import ABCMeta 
+        class Test(object):
+            __metaclass__ = ABCMeta
+        """))
+        klass = astroid['Test']
+
+        metaclass = klass.metaclass()
+        self.assertIsInstance(metaclass, scoped_nodes.Class)
+        self.assertEqual(metaclass.name, 'ABCMeta')
+
+    @require_version('2.7')
+    def test_newstyle_and_metaclass_good(self):
+        astroid = abuilder.string_build(dedent("""
+        from abc import ABCMeta 
+        class Test:
+            __metaclass__ = ABCMeta
+        """))
+        klass = astroid['Test']
+        self.assertTrue(klass.newstyle)
+
+    def test_newstyle_and_metaclass_bad(self):
+        astroid = abuilder.string_build(dedent("""
+        class Test:
+            __metaclass__ = int
+        """))
+        klass = astroid['Test']
+        self.assertFalse(klass.newstyle)
+
+    @require_version('2.7')
+    def test_parent_metaclass(self):
+        astroid = abuilder.string_build(dedent("""
+        from abc import ABCMeta
+        class Test:
+            __metaclass__ = ABCMeta
+        class SubTest(Test): pass
+        """))
+        klass = astroid['SubTest']
+        self.assertTrue(klass.newstyle)
+        metaclass = klass.metaclass()
+        self.assertIsInstance(metaclass, scoped_nodes.Class)
+        self.assertEqual(metaclass.name, 'ABCMeta')
+
 
 __all__ = ('ModuleNodeTC', 'ImportNodeTC', 'FunctionNodeTC', 'ClassNodeTC')