Optimize at method. Add more tests.

redbaron/base_nodes.py

@@ -213,15 +213,16 @@ class GenericNodesUtils(object):
             return None
     def at(self, line_no):
-        if line_no > self.absolute_bounding_box.bottom_right.line or \
-                        line_no < self.absolute_bounding_box.top_left.line:
+        if not self.absolute_bounding_box.top_left.line <= \
+                line_no <= self.absolute_bounding_box.bottom_right.line:
             raise IndexError("Line number {0} is outside of the file".format(line_no))
         path = Path.from_baron_path(self, baron.path.position_to_path(self.fst(), (line_no, 1)))
         node = path.node if path else None
         if node is not None and hasattr(node, 'type') and node.type == 'endl':
-            # I am not sure that it is the best solution from design point,
-            # but it is the first idea which is come to mind
-            return node.next_recursive
+            if node.next:
+                return list(self._iter_in_rendering_order(node.next))[1]
+            elif node.next_recursive:
+                return list(self._iter_in_rendering_order(node.next_recursive))[1]
         return node
     def _string_to_node_list(self, string, parent, on_attribute):

tests/test_at.py

@@ -21,8 +21,8 @@ def test_at_def():
     assert fst_def.at(1) is fst_def.find('DecoratorNode')
     assert fst_def.at(2) is fst_def
     assert fst_def.at(3) is fst_def.def_
-    assert fst_def.at(4) is fst_def.find_all('AssignmentNode')[0]
-    assert fst_def.at(5) is fst_def.find_all('AssignmentNode')[1]
+    assert fst_def.at(4) is fst_def.find('NameNode', value='c')
+    assert fst_def.at(5) is fst_def.find('NameNode', value='e')
 fst_class = RedBaron("""\
@@ -37,6 +37,19 @@ class Foo(object):
 def test_at_class():
     assert fst_class.at(1) is fst_class.class_
     assert fst_class.at(2) is fst_class.find_all('DefNode')[0]
-    assert fst_class.at(3) is fst_class.find('AssignmentNode')
+    assert fst_class.at(3) is fst_class.find('AtomtrailersNode')
     assert fst_class.at(4) is fst_class.find_all('DefNode')[1]
     assert fst_class.at(5) is fst_class.find('ReturnNode')
+
+
+fst_simple = RedBaron("""\
+a = 5
+
+b = 6
+""")
+
+
+def test_at_simple():
+    assert fst_simple.at(1) is fst_simple.find_all('NameNode')[0]
+    assert fst_simple.at(2) is fst_simple
+    assert fst_simple.at(3) is fst_simple.find_all('NameNode')[1]