Fix bugs in at method. Update test case

redbaron/base_nodes.py

@@ -210,12 +210,23 @@ class GenericNodesUtils(object):
         return path.node if path else None
     def at(self, line_no):
-        if not self.absolute_bounding_box.top_left.line <= \
-                line_no <= self.absolute_bounding_box.bottom_right.line:
+        if not 0 <= line_no <= self.absolute_bounding_box.bottom_right.line:
             raise IndexError("Line number {0} is outside of the file".format(line_no))
         node = self.find_by_position((line_no, 1))
-        if node is not None and hasattr(node, 'type') and node.type == 'endl':
-            return list(self._iter_in_rendering_order(node.next_recursive))[1]
+        if node.absolute_bounding_box.top_left.line == line_no:
+            if hasattr(node.parent, 'absolute_bounding_box') and \
+                            node.parent.absolute_bounding_box.top_left.line == line_no and \
+                    node.parent.parent is not None:
+                return node.parent
+            return node
+        elif node is not None and hasattr(node, 'next_rendered'):
+            return list(self._iter_in_rendering_order(node.next_rendered))[0]
+        elif node.parent is None:
+            node = node.data[0][0]
+            while True:
+                if node.absolute_bounding_box.top_left.line == line_no:
+                    return node
+                node = node.next_rendered
         return node
     def _string_to_node_list(self, string, parent, on_attribute):

tests/test_at.py

@@ -2,7 +2,6 @@
 """Test for at method"""
-import pytest
 import redbaron
 from redbaron import RedBaron
@@ -15,13 +14,19 @@ class Foo(object):
     def bar(self):
         for x in range(5):
             yield self.a + x
+
+setup(name='redbaron',
+      version='0.6.1')
 """)
 def test_at():
     assert red.at(1) is red.class_
     assert red.at(2) is red.find_all('DefNode')[0]
-    assert red.at(3) is red.find('AtomtrailersNode')
+    assert red.at(3) is red.find('AssignmentNode')
     assert red.at(4) is red.find_all('DefNode')[1]
     assert red.at(5) is red.find('ForNode')
     assert red.at(6) is red.find('YieldNode')
+    assert red.at(7) is red.find_all('EndlNode')[6]
+    assert red.at(8) is red.find_all('AtomTrailersNode')[3]
+    assert red.at(9) is red.find_all('CallArgumentNode')[2]