test and fix (no crash at least) block_range methods

2 files changed, 72 insertions, 33 deletions

diff --git a/nodes.py b/nodes.py
index 0878136..326e6db 100644
--- a/nodes.py
+++ b/nodes.py
@@ -390,17 +390,12 @@ Module.block_range = object_block_range
 
 # XXX only if compiler mode ?
 def if_block_range(node, lineno):
-    """handle block line numbers range for if/elif statements
-    """
-    last = None
-    for test, testbody in node.tests[1:]:
-        if lineno == testbody.source_line():
-            return lineno, lineno
-        if lineno <= testbody.last_source_line():
-            return lineno, testbody.last_source_line()
-        if last is None:
-            last = testbody.source_line() - 1
-    return elsed_block_range(node, lineno, last)
+    """handle block line numbers range for if/elif statements"""
+    if lineno == node.body[0].source_line():
+        return lineno, lineno
+    if lineno <= node.body[-1].last_source_line():
+        return lineno, node.body[-1].last_source_line()
+    return elsed_block_range(node, lineno, node.body[0].source_line() - 1)
 
 If.block_range = if_block_range
 
@@ -408,13 +403,13 @@ def try_except_block_range(node, lineno):
     """handle block line numbers range for try/except statements
     """
     last = None
-    for excls, exinst, exbody in node.handlers:
-        if excls and lineno == excls.source_line():
+    for exhandler in node.handlers:
+        if exhandler.type and lineno == exhandler.type.source_line():
             return lineno, lineno
-        if exbody.source_line() <= lineno <= exbody.last_source_line():
-            return lineno, exbody.last_source_line()
+        if exhandler.body[0].source_line() <= lineno <= exhandler.body[-1].last_source_line():
+            return lineno, exhandler.body[-1].last_source_line()
         if last is None:
-            last = exbody.source_line() - 1
+            last = exhandler.body[0].source_line() - 1
     return elsed_block_range(node, lineno, last)
 
 TryExcept.block_range = try_except_block_range
@@ -425,10 +420,10 @@ def elsed_block_range(node, lineno, last=None):
     """
     if lineno == node.source_line():
         return lineno, lineno
-    if node.else_:
-        if lineno >= node.else_.source_line():
-            return lineno, node.else_.last_source_line()
-        return lineno, node.else_.source_line() - 1
+    if node.orelse:
+        if lineno >= node.orelse[0].source_line():
+            return lineno, node.orelse[-1].last_source_line()
+        return lineno, node.orelse[0].source_line() - 1
     return lineno, last or node.last_source_line()
 
 TryFinally.block_range = elsed_block_range
diff --git a/test/unittest_nodes.py b/test/unittest_nodes.py
index f78cbcc..7c2a95b 100644
--- a/test/unittest_nodes.py
+++ b/test/unittest_nodes.py
@@ -21,7 +21,22 @@ from data import module as test_module
 
 abuilder = builder.ASTNGBuilder()
 
-IF_CODE = """
+class _NodeTC(testlib.TestCase):
+    """test transformation of If Node"""
+    CODE = None
+    @property
+    def astng(self):
+        try:
+            return self.__class__.__dict__['CODE_ASTNG']
+        except KeyError:
+            astng = abuilder.string_build(self.CODE)
+            self.__class__.CODE_ASTNG = astng
+            return astng
+
+
+class IfNodeTC(_NodeTC):
+    """test transformation of If Node"""
+    CODE = """
 if 0:
     print
 
@@ -43,21 +58,50 @@ elif func():
     pass
 else:
     raise
-"""
-
-class IfNodeTC(testlib.TestCase):
-    """test transformation of If Node"""
-
+    """
+        
     def test_if_elif_else_node(self):
         """test transformation for If node"""
-        module = abuilder.string_build(IF_CODE)
-        self.assertEquals(len(module.body), 4)
-        for stmt in module.body:
+        self.assertEquals(len(self.astng.body), 4)
+        for stmt in self.astng.body:
             self.assertIsInstance( stmt, nodes.If)
-        self.assert_(not module.body[0].orelse) # simple If
-        self.assertIsInstance(module.body[1].orelse[0], nodes.Pass) # If / else
-        self.assertIsInstance(module.body[2].orelse[0], nodes.If) # If / elif
-        self.assertIsInstance(module.body[3].orelse[0].orelse[0], nodes.If)
+        self.failIf(self.astng.body[0].orelse) # simple If
+        self.assertIsInstance(self.astng.body[1].orelse[0], nodes.Pass) # If / else
+        self.assertIsInstance(self.astng.body[2].orelse[0], nodes.If) # If / elif
+        self.assertIsInstance(self.astng.body[3].orelse[0].orelse[0], nodes.If)
+
+    def test_block_range(self):
+        # XXX ensure expected values
+        self.assertEquals(self.astng.block_range(1), (0, 22))
+        self.assertEquals(self.astng.block_range(10), (0, 22)) # XXX (10, 22) ?
+        self.assertEquals(self.astng.body[1].block_range(5), (5, 6))
+        self.assertEquals(self.astng.body[1].block_range(6), (6, 6))
+        self.assertEquals(self.astng.body[1].orelse[0].block_range(0), (0, 8))
+        self.assertEquals(self.astng.body[1].orelse[0].block_range(7), (7, 8))
+        self.assertEquals(self.astng.body[1].orelse[0].block_range(8), (8, 8))
+
+class TryExceptNodeTC(_NodeTC):
+    CODE = """
+try:
+    print 'pouet'
+except IOError:
+    pass
+except UnicodeError:
+    print
+else:
+    print
+    """
+    def test_block_range(self):
+        # XXX ensure expected values
+        self.assertEquals(self.astng.body[0].block_range(0), (0, 8))
+        self.assertEquals(self.astng.body[0].block_range(1), (1, 8))
+        self.assertEquals(self.astng.body[0].block_range(2), (2, 2))
+        self.assertEquals(self.astng.body[0].block_range(3), (3, 8))
+        self.assertEquals(self.astng.body[0].block_range(4), (4, 4))
+        self.assertEquals(self.astng.body[0].block_range(5), (5, 5))
+        self.assertEquals(self.astng.body[0].block_range(6), (6, 6))
+        self.assertEquals(self.astng.body[0].block_range(7), (7, 7))
+        self.assertEquals(self.astng.body[0].block_range(8), (8, 8))
         
 MODULE = abuilder.module_build(test_module)
 MODULE2 = abuilder.file_build('data/module2.py', 'data.module2')