[enh] new helper: .insert_after

PyCQA · Jan 31, 2015 · ea5ed76 · ea5ed76
1 parent 8ad0dbd
commit ea5ed76
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/redbaron.py b/redbaron.py
@@ -821,6 +821,7 @@ def _get_helpers(self):
             'index_on_parent',
             'index_on_parent_raw',
             'insert_before',
+            'insert_after',
         ])
         return [x for x in dir(self) if not x.startswith("_") and x not in not_helpers and inspect.ismethod(getattr(self, x))]
 
@@ -1013,6 +1014,10 @@ def insert_before(self, value):
         self.parent.insert(self.index_on_parent, value)
 
 
+    def insert_after(self, value):
+        self.parent.insert(self.index_on_parent + 1, value)
+
+
 class CodeBlockNode(Node):
     def _string_to_node_list(self, string, parent, on_attribute):
         if on_attribute == "value":

diff --git a/tests/test_initial_parsing.py b/tests/test_initial_parsing.py
@@ -551,3 +551,9 @@ def test_insert_before():
     red = RedBaron("a = 1\nprint pouet\n")
     red.print_.insert_before("chocolat")
     assert red.dumps() == "a = 1\nchocolat\nprint pouet\n"
+
+
+def test_insert_after():
+    red = RedBaron("a = 1\nprint pouet\n")
+    red.print_.insert_after("chocolat")
+    assert red.dumps() == "a = 1\nprint pouet\nchocolat\n"