Proper metaclass discovery when encountering __metaclass__.

1 files changed, 8 insertions, 4 deletions

diff --git a/rebuilder.py b/rebuilder.py
index 862de85..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'
+            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