Merge branch 'pull/115'

CHANGELOG

@@ -6,6 +6,7 @@ Changelog
 
 - fix help() after append
 - fix _synchronise() for base_nodes to avoid recursion in __repr__ function if code run not in the shell
+- add at method
 
 0.6.2 (2016-10-03)
 ----------------

TODO.md

@@ -2,9 +2,7 @@
 
 ### Important
 
-- raise an exception on .find/.find_all if the identifier given doesn't exists
 - .help() seems really slow on big piece of code (for example RedBaron("baron/grammator.py").read())("dict")[0].help() is suuuuuuuuuuuuuuuuper slow)
-- .at() return the first item starting at line X
 - .rename() (name -> value, def/class -> name)
 - .replace() expect a whole valid python program. This could be fixed by look at "on_attribute" and resetting itself like that.
 

docs/other.rst

@@ -476,6 +476,17 @@ You can find which node is located at a given line and column:
     red.find_by_position((1, 5))
     red.find_by_position((1, 6)) # '(' is not a redbaron node
 
+.at()
+-------------------
+
+Returns first node at specific line
+
+.. ipython:: python
+
+    red = RedBaron("def a():\n return 42")
+    red.at(1) # Gives DefNode
+    red.at(2) # Gives ReturnNode
+
 .. _Node.from_fst:
 
 Node.from_fst()

redbaron/base_nodes.py

@@ -211,10 +211,27 @@ def absolute_bounding_box(self):
 
     def find_by_position(self, position):
         path = Path.from_baron_path(self, baron.path.position_to_path(self.fst(), position))
-        if path:
-            return path.node
-        else:
-            return None
+        return path.node if path else None
+
+    def at(self, line_no):
+        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.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):
         return NodeList.from_fst(baron.parse(string), parent=parent, on_attribute=on_attribute)
@@ -880,6 +897,7 @@ def _get_helpers(self):
             'findAll',
             'find_by_path',
             'find_by_position',
+            'at',
             'from_fst',
             'fst',
             'generate_identifiers',

tests/test_at.py

@@ -0,0 +1,32 @@
+# -*- coding:Utf-8 -*-
+
+"""Test for at method"""
+
+import redbaron
+from redbaron import RedBaron
+
+redbaron.DEBUG = True
+
+red = RedBaron("""\
+class Foo(object):
+    def __init__(self):
+        self.a = None
+    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('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]