patch transformer to extract correct line information

6 files changed, 251 insertions, 7 deletions

diff --git a/ChangeLog b/ChangeLog
index d102a0d..43c1902 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -9,6 +9,7 @@ Change log for the astng package
       still needed though)
     * drop python2.2 support: remove code that was working around python2.2
     * fixed generator expression scope bug
+    * patch transformer to extract correct line information
 	
 	
 2006-04-19  --  0.16.0
diff --git a/builder.py b/builder.py
index cc6df36..5ee282b 100644
--- a/builder.py
+++ b/builder.py
@@ -28,8 +28,7 @@ TODO:
 :contact:   mailto:thenault@gmail.com
 """
 
-__revision__ = "$Id: builder.py,v 1.54 2006-03-14 15:21:33 syt Exp $"
-__doctype__ = "restructuredtext en"
+__docformat__ = "restructuredtext en"
 
 import sys
 from os.path import splitext, basename, dirname, exists, abspath
@@ -47,6 +46,107 @@ from logilab.astng.utils import ASTWalker
 from logilab.astng._exceptions import ASTNGBuildingException
 from logilab.astng.raw_building import *
 
+import token
+from compiler import transformer, consts
+from types import TupleType
+from logilab.common.astutils import cvrtr
+
+def fromto_lineno(asttuple):
+    """return the minimum and maximum line number of the given ast tuple"""
+    return from_lineno(asttuple), to_lineno(asttuple)
+def from_lineno(asttuple):
+    """return the minimum line number of the given ast tuple"""
+    if type(asttuple[1]) is TupleType:
+        return from_lineno(asttuple[1])
+    return asttuple[2]
+def to_lineno(asttuple):
+    """return the maximum line number of the given ast tuple"""
+    if type(asttuple[-1]) is TupleType:
+        return to_lineno(asttuple[-1])
+    return asttuple[2]
+
+def fix_lineno(node, fromast, toast=None):
+    if hasattr(node, 'fromlineno'):
+        return node    
+    #print 'fixing', id(node), id(node.__dict__), node.__dict__.keys(), repr(node)
+    if isinstance(node, nodes.Stmt):
+        node.fromlineno = from_lineno(fromast)#node.nodes[0].fromlineno
+        node.tolineno = node.nodes[-1].tolineno
+        return node
+    if toast is None:
+        node.fromlineno, node.tolineno = fromto_lineno(fromast)
+    else:
+        node.fromlineno, node.tolineno = from_lineno(fromast), to_lineno(toast)
+    #print 'fixed', id(node)
+    return node
+
+BaseTransformer = transformer.Transformer
+
+COORD_MAP = {
+    # if: test ':' suite ('elif' test ':' suite)* ['else' ':' suite]
+    'if': (0, 0),
+    # 'while' test ':' suite ['else' ':' suite]
+    'while': (0, 1),
+    # 'for' exprlist 'in' exprlist ':' suite ['else' ':' suite]
+    'for': (0, 3),
+    # 'try' ':' suite (except_clause ':' suite)+ ['else' ':' suite]
+    'try': (0, 0),
+    # | 'try' ':' suite 'finally' ':' suite
+    
+    }
+
+def fixlineno_wrap(function, stype):
+    def fixlineno_wrapper(self, nodelist):
+        node = function(self, nodelist)            
+        idx1, idx2 = COORD_MAP.get(stype, (0, -1))
+        return fix_lineno(node, nodelist[idx1], nodelist[idx2])
+    return fixlineno_wrapper
+nodes.Module.fromlineno = 0
+nodes.Module.tolineno = 0
+class ASTNGTransformer(BaseTransformer):
+    """ovverides transformer for a better source line number handling"""
+    def com_NEWLINE(self, *args):
+        # A ';' at the end of a line can make a NEWLINE token appear
+        # here, Render it harmless. (genc discards ('discard',
+        # ('const', xxxx)) Nodes)
+        lineno = args[0][1]
+        # don't put fromlineno/tolineno on Const None to mark it as dynamically
+        # added, without "physical" reference in the source
+        n = nodes.Discard(nodes.Const(None))
+        n.fromlineno = n.tolineno = lineno
+        return n    
+    def com_node(self, node):
+        res = self._dispatch[node[0]](node[1:])
+        return fix_lineno(res, node)
+    def com_assign(self, node, assigning):
+        res = BaseTransformer.com_assign(self, node, assigning)
+        return fix_lineno(res, node)
+    def com_apply_trailer(self, primaryNode, nodelist):
+        node = BaseTransformer.com_apply_trailer(self, primaryNode, nodelist)
+        return fix_lineno(node, nodelist)
+    
+##     def atom(self, nodelist):
+##         node = BaseTransformer.atom(self, nodelist)
+##         return fix_lineno(node, nodelist[0], nodelist[-1])
+    
+    def funcdef(self, nodelist):
+        node = BaseTransformer.funcdef(self, nodelist)
+        # XXX decorators
+        return fix_lineno(node, nodelist[-5], nodelist[-3])
+    def classdef(self, nodelist):
+        node = BaseTransformer.classdef(self, nodelist)
+        return fix_lineno(node, nodelist[0], nodelist[-2])
+            
+# wrap *_stmt methods
+for name in dir(BaseTransformer):
+    if name.endswith('_stmt') and not (name in ('com_stmt',
+                                                'com_append_stmt')
+                                       or name in ASTNGTransformer.__dict__):
+        setattr(BaseTransformer, name,
+                fixlineno_wrap(getattr(BaseTransformer, name), name[:-5]))
+            
+transformer.Transformer = ASTNGTransformer
+
 # ast NG builder ##############################################################
 
 class ASTNGBuilder:
diff --git a/nodes.py b/nodes.py
index fc17df3..380a425 100644
--- a/nodes.py
+++ b/nodes.py
@@ -237,7 +237,6 @@ class NodeNG:
             return name
         return None
     
-
 extend_class(Node, NodeNG)
 
 # block range overrides #######################################################
diff --git a/test/data/format.py b/test/data/format.py
new file mode 100644
index 0000000..7379706
--- /dev/null
+++ b/test/data/format.py
@@ -0,0 +1,34 @@
+"""A multiline string
+"""
+
+function('aeozrijz\
+earzer', hop)
+# XXX write test
+x = [i for i in range(5)
+     if i % 4]
+
+fonction(1,
+         2,
+         3,
+         4)
+
+def definition(a,
+               b,
+               c):
+    return a + b + c
+
+class debile(dict,
+             object):
+    pass
+
+if aaaa: pass
+else:
+    aaaa,bbbb = 1,2
+    aaaa,bbbb = bbbb,aaaa
+# XXX write test
+hop = \
+    aaaa
+
+
+__revision__.lower();
+
diff --git a/test/unittest_builder.py b/test/unittest_builder.py
index 41a29b9..511402a 100644
--- a/test/unittest_builder.py
+++ b/test/unittest_builder.py
@@ -12,7 +12,7 @@
 # 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 """tests for the astng builder module
 
-Copyright (c) 2003-2005 LOGILAB S.A. (Paris, FRANCE).
+Copyright (c) 2003-2006 LOGILAB S.A. (Paris, FRANCE).
 http://www.logilab.fr/ -- mailto:contact@logilab.fr
 """
 
@@ -21,14 +21,124 @@ import sys
 from os.path import join, abspath
 
 from logilab.common.testlib import TestCase, unittest_main
+from logilab.common.astutils import cvrtr
 from unittest_inference import get_name_node
-
+from pprint import pprint
+        
 from logilab.astng import builder, nodes, Module, YES
 
 import data
 from data import module as test_module
 
+class TransformerTC(TestCase):
+        
+    def setUp(self):
+        transformer = builder.ASTNGTransformer()
+        self.astng = transformer.parsesuite(open('data/format.py').read())
+
+    def test_callfunc_lineno(self):
+        stmts = self.astng.getChildNodes()[0].nodes
+        # on line 4:
+        #    function('aeozrijz\
+        #    earzer', hop)
+        discard = stmts[0]
+        self.assertIsInstance(discard, nodes.Discard)
+        self.assertEquals(discard.fromlineno, 4)
+        self.assertEquals(discard.tolineno, 5)
+        callfunc = discard.expr
+        self.assertIsInstance(callfunc, nodes.CallFunc)
+        self.assertEquals(callfunc.fromlineno, 4)
+        self.assertEquals(callfunc.tolineno, 5)
+        name = callfunc.node
+        self.assertIsInstance(name, nodes.Name)
+        self.assertEquals(name.fromlineno, 4)
+        self.assertEquals(name.tolineno, 4)
+        strarg = callfunc.args[0]
+        self.assertIsInstance(strarg, nodes.Const)
+        self.assertEquals(strarg.fromlineno, 5) # no way for this one (is 4 actually)
+        self.assertEquals(strarg.tolineno, 5)
+        namearg = callfunc.args[1]
+        self.assertIsInstance(namearg, nodes.Name)
+        self.assertEquals(namearg.fromlineno, 5)
+        self.assertEquals(namearg.tolineno, 5)
+        # on line 10:
+        #    fonction(1,
+        #             2,
+        #             3,
+        #             4)
+        discard = stmts[2]
+        self.assertIsInstance(discard, nodes.Discard)
+        self.assertEquals(discard.fromlineno, 10)
+        self.assertEquals(discard.tolineno, 13)
+        callfunc = discard.expr
+        self.assertIsInstance(callfunc, nodes.CallFunc)
+        self.assertEquals(callfunc.fromlineno, 10)
+        self.assertEquals(callfunc.tolineno, 13)
+        name = callfunc.node
+        self.assertIsInstance(name, nodes.Name)
+        self.assertEquals(name.fromlineno, 10)
+        self.assertEquals(name.tolineno, 10)
+        for i, arg in enumerate(callfunc.args):
+            self.assertIsInstance(arg, nodes.Const)
+            self.assertEquals(arg.fromlineno, 10+i)
+            self.assertEquals(arg.tolineno, 10+i)
+
+    def test_function_lineno(self):
+        stmts = self.astng.getChildNodes()[0].nodes
+        # on line 15:
+        #    def definition(a,
+        #                   b,
+        #                   c):
+        #        return a + b + c
+        function = stmts[3]
+        self.assertIsInstance(function, nodes.Function)
+        self.assertEquals(function.fromlineno, 15)
+        self.assertEquals(function.tolineno, 17)
+        code = function.code
+        self.assertIsInstance(code, nodes.Stmt)
+##         self.assertEquals(code.fromlineno, 18)
+##         self.assertEquals(code.tolineno, 18)
+        return_ = code.nodes[0]
+        self.assertIsInstance(return_, nodes.Return)
+        self.assertEquals(return_.fromlineno, 18)
+        self.assertEquals(return_.tolineno, 18)
+
+    def test_class_lineno(self):
+        stmts = self.astng.getChildNodes()[0].nodes
+        # on line 20:
+        #    class debile(dict,
+        #                 object):
+        #       pass
+        class_ = stmts[4]
+        self.assertIsInstance(class_, nodes.Class)
+        self.assertEquals(class_.fromlineno, 20)
+        self.assertEquals(class_.tolineno, 21)
+        code = class_.code
+        self.assertIsInstance(code, nodes.Stmt)
+##         self.assertEquals(code.fromlineno, 18)
+##         self.assertEquals(code.tolineno, 18)
+        pass_ = code.nodes[0]
+        self.assertIsInstance(pass_, nodes.Pass)
+        self.assertEquals(pass_.fromlineno, 22)
+        self.assertEquals(pass_.tolineno, 22)
+
+    def test_if_lineno(self):
+        stmts = self.astng.getChildNodes()[0].nodes
+        # on line 20:
+        #    if aaaa: pass
+        #    else:
+        #        aaaa,bbbb = 1,2
+        #        aaaa,bbbb = bbbb,aaaa
+        if_ = stmts[5]
+        self.assertIsInstance(if_, nodes.If)
+        self.assertEquals(if_.fromlineno, 24)
+        self.assertEquals(if_.tolineno, 24)
+        else_ = if_.else_
+        self.assertIsInstance(else_, nodes.Stmt)
+        self.assertEquals(else_.fromlineno, 25)
+        self.assertEquals(else_.tolineno, 27)
 
+        
 class BuilderTC(TestCase):
         
     def setUp(self):
diff --git a/test/unittest_manager.py b/test/unittest_manager.py
index fd01036..ed42b87 100644
--- a/test/unittest_manager.py
+++ b/test/unittest_manager.py
@@ -34,9 +34,9 @@ class ASTNGManagerTC(unittest.TestCase):
         
     def test_package_node(self):
         obj = self.manager.from_directory('data')
-        expected_short = ['SSL1', '__init__', 'all', 'appl', 'module', 'module2',
+        expected_short = ['SSL1', '__init__', 'all', 'appl', 'format', 'module', 'module2',
                           'noendingnewline', 'nonregr', 'notall']
-        expected_long = ['SSL1', 'data', 'data.all', 'appl', 'data.module',
+        expected_long = ['SSL1', 'data', 'data.all', 'appl', 'data.format', 'data.module',
                          'data.module2', 'data.noendingnewline', 'data.nonregr',
                          'data.notall']
         self.assertEquals(obj.keys(), expected_short)