Add support for metaclass for Python 3.

2 files changed, 20 insertions, 2 deletions

diff --git a/rebuilder.py b/rebuilder.py
index 7f4c7a71..fbc7ba8a 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,
@@ -99,6 +99,13 @@ def _init_set_doc(node, newnode):
     except IndexError:
         pass # ast built from scratch
 
+def _infer_metaclass(node):    
+    if isinstance(node, Name):
+        return node.id
+    elif isinstance(node, Attribute):
+        # TODO: should we retrieve the fully qualified name of the class?
+        return node.attr
+
 def _lineno_parent(oldnode, newnode, parent):
     newnode.parent = parent
     if hasattr(oldnode, 'lineno'):
@@ -934,6 +941,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 = _infer_metaclass(keyword.value)
+                break
+        return newnode
+
 if sys.version_info >= (3, 0):
     TreeRebuilder = TreeRebuilder3k
 
diff --git a/scoped_nodes.py b/scoped_nodes.py
index 524e3f68..f9df97d9 100644
--- a/scoped_nodes.py
+++ b/scoped_nodes.py
@@ -991,9 +991,12 @@ class Class(Statement, LocalsDictNodeNG, FilterStmtsMixin):
         if missing:
             raise InferenceError()
 
+    _metaclass = None
     def metaclass(self):
         """ return the metaclass of this class """
-        # TODO: handle class A(metaclass=B) for Python 3
+        if self._metaclass:
+            return self._metaclass
+
         try:
             meta = Instance(self).getattr('__metaclass__')[0]
         except NotFoundError: