Make the handling of inserted newlines more consistent

redbaron/base_nodes.py

@@ -1114,10 +1114,12 @@ def parse_code_block(self, string, parent, on_attribute):
 
         # putting this in the string template will fail, need at least some indent
         if indentation == 0:
-            clean_string = "    " + "\n    ".join(clean_string.split("\n"))
-            clean_string = clean_string.rstrip()
+            if clean_string == '':
+                clean_string = '    '
+            else:
+                clean_string = "\n".join(["    " + x if x else '' for x in clean_string.split("\n")])
 
-        fst = baron.parse("def a():\n%s\n" % clean_string)[0]["value"]
+        fst = baron.parse("def a():\n%s" % clean_string)[0]["value"]
 
         result = NodeList.from_fst(fst, parent=parent, on_attribute=on_attribute)
 
@@ -1128,28 +1130,7 @@ def parse_code_block(self, string, parent, on_attribute):
         elif indentation < target_indentation:
             result.increase_indentation(target_indentation - indentation)
 
-        endl_base_node = Node.from_fst({'formatting': [], 'indent': '', 'type': 'endl', 'value': '\n'},
-                                       on_attribute=on_attribute, parent=parent)
-
-        if (self.on_attribute == "root" and self.next) or (not self.next and self.parent and self.parent.next):
-            # I need to finish with 3 endl nodes
-            if not all(map(lambda x: x.type == "endl", result[-1:])):
-                result.append(endl_base_node.copy())
-            elif not all(map(lambda x: x.type == "endl", result[-2:])):
-                result.append(endl_base_node.copy())
-                result.append(endl_base_node.copy())
-            elif not all(map(lambda x: x.type == "endl", result[-3:])):
-                result.append(endl_base_node.copy())
-                result.append(endl_base_node.copy())
-                result.append(endl_base_node.copy())
-        elif self.next:
-            # I need to finish with 2 endl nodes
-            if not all(map(lambda x: x.type == "endl", result[-2:])):
-                result.append(endl_base_node.copy())
-            elif not all(map(lambda x: x.type == "endl", result[-3:])):
-                result.append(endl_base_node.copy())
-                result.append(endl_base_node.copy())
-
+        if self.next:
             result[-1].indent = self.indentation
 
         return result
@@ -1721,8 +1702,7 @@ def modify_last_indentation(node, indentation):
             log("[%s] %s", position, i)
 
             if might_need_separator and i[0].type != "endl" and (
-                        not previous or previous.type != "endl") and not isinstance(previous, (
-                    CodeBlockNode, redbaron.nodes.IfelseblockNode)):
+                        not previous or previous.type != "endl") and not isinstance(previous, CodeBlockNode):
                 log(">> Previous line has content and current needs to be indented, append separator to indent it")
                 expected_list.append(generate_separator())
                 log("-- current result: %s", ["".join(map(lambda x: x.dumps(), expected_list))])

redbaron/nodes.py

@@ -708,7 +708,7 @@ def _string_to_node(self, string, parent, on_attribute):
             raise Exception("Unhandled case")
 
 
-class IfelseblockNode(Node):
+class IfelseblockNode(CodeBlockNode):
     def _string_to_node_list(self, string, parent, on_attribute):
         if on_attribute != "value":
             return super(IfelseblockNode, self)._string_to_node_list(string, parent=parent, on_attribute=on_attribute)

tests/test_insert.py

@@ -318,6 +318,23 @@ def a(self, a):
 """)
 
 
+def test_insert_line_in_def_with_trailing_if():
+    red = RedBaron("""\
+def a(self, a):
+    if a == 42:
+        return True
+""")
+
+    red.def_.insert(0, "a = 1")
+
+    assert_with_indent(red, """\
+def a(self, a):
+    a = 1
+    if a == 42:
+        return True
+""")
+
+
 def test_insert_nested_line_in_def_with_if():
     red = RedBaron("""\
 def a(self, a):

tests/test_setter.py

@@ -225,72 +225,130 @@ def d():
 
 def test_set_attr_def_advanced_dont_break_next_block_indent():
     red = RedBaron(code_for_block_setattr)
-    red.find("def", name="c").value = "return 42"
-    assert len(red.find("def", name="c")("endl")) == 4
-    assert red.find("def", name="c").value.node_list[-1].indent == ""
+    c = red.find("def", name="c")
+    c.value = "return 42"
+    assert c.dumps() == """\
+def c():
+    return 42
+"""
+    assert len(c("endl")) == 2
+    assert c.value.node_list[-1].indent == ""
 
 
 def test_set_attr_def_advanced_dont_break_next_block_indent_one_endl():
     red = RedBaron(code_for_block_setattr)
-    red.find("def", name="c").value = "return 42\n"
-    assert len(red.find("def", name="c")("endl")) == 4
-    assert red.find("def", name="c").value.node_list[-1].indent == ""
+    c = red.find("def", name="c")
+    c.value = "return 42\n"
+    assert c.dumps() == """\
+def c():
+    return 42
+"""
+    assert len(c("endl")) == 2
+    assert c.value.node_list[-1].indent == ""
 
 
 def test_set_attr_def_advanced_dont_break_next_block_indent_two_endl():
     red = RedBaron(code_for_block_setattr)
+    c = red.find("def", name="c")
+    assert len(c("endl")) == 6
     red.find("def", name="c").value = "return 42\n\n"
-    assert len(red.find("def", name="c")("endl")) == 4
+    assert c.dumps() == """\
+def c():
+    return 42
+
+"""
+    assert len(red.find("def", name="c")("endl")) == 3
     assert red.find("def", name="c").value.node_list[-1].indent == ""
 
 
 def test_set_attr_def_advanced_in_class_dont_break_next_block_indent():
     red = RedBaron(code_for_block_setattr)
-    red.find("def", name="a").value = "return 42"
-    assert len(red.find("def", name="a")("endl")) == 3
-    assert red.find("def", name="a").value.node_list[-1].indent == "    "
+    a = red.find("def", name="a")
+    assert a.dumps() == """\
+def a():
+        pass
+
+    """
+    a.value = "return 42"
+    assert a.dumps() == """\
+def a():
+        return 42
+    """
+    assert len(a("endl")) == 2
+    assert a.value.node_list[-1].indent == "    "
 
 
 def test_set_attr_def_advanced_in_class_dont_break_next_block_indent_one_endl():
     red = RedBaron(code_for_block_setattr)
-    red.find("def", name="a").value = "return 42\n"
-    assert len(red.find("def", name="a")("endl")) == 3
-    assert red.find("def", name="a").value.node_list[-1].indent == "    "
+    a = red.find("def", name="a")
+    a.value = "return 42\n"
+    assert a.dumps() == """\
+def a():
+        return 42
+    """
+    assert len(a("endl")) == 2
+    assert a.value.node_list[-1].indent == "    "
 
 
 def test_set_attr_def_advanced_in_class_at_the_end_dont_break_next_block_indent():
     red = RedBaron(code_for_block_setattr)
-    red.find("def", name="b").value = "return 42"
-    assert len(red.find("def", name="b")("endl")) == 4
-    assert red.find("def", name="b").value.node_list[-1].indent == ""
+    b = red.find("def", name="b")
+    b.value = "return 42"
+    assert b.dumps() == """\
+def b():
+        return 42
+"""
+    assert len(b("endl")) == 2
+    assert b.value.node_list[-1].indent == ""
 
 
 def test_set_attr_def_advanced_in_class_at_the_end_dont_break_next_block_indent_one_endl():
     red = RedBaron(code_for_block_setattr)
-    red.find("def", name="b").value = "return 42\n"
-    assert len(red.find("def", name="b")("endl")) == 4
-    assert red.find("def", name="b").value.node_list[-1].indent == ""
+    b = red.find("def", name="b")
+    b.value = "return 42\n"
+    assert b.dumps() == """\
+def b():
+        return 42
+"""
+    assert len(b("endl")) == 2
+    assert b.value.node_list[-1].indent == ""
 
 
 def test_set_attr_def_advanced_in_class_at_the_end_dont_break_next_block_indent_two_endl():
     red = RedBaron(code_for_block_setattr)
-    red.find("def", name="b").value = "return 42\n\n"
-    assert len(red.find("def", name="b")("endl")) == 4
-    assert red.find("def", name="b").value.node_list[-1].indent == ""
+    b = red.find("def", name="b")
+    b.value = "return 42\n\n"
+    assert b.dumps() == """\
+def b():
+        return 42
+
+"""
+    assert len(b("endl")) == 3
+    assert b.value.node_list[-1].indent == ""
 
 
 def test_set_attr_def_advanced_inline_dont_break_next_block_indent():
     red = RedBaron(code_for_block_setattr)
-    red.find("def", name="zomg").value = "return 42"
-    assert len(red.find("def", name="zomg")("endl")) == 3
-    assert red.find("def", name="zomg").value.node_list[-1].indent == "    "
+    zomg = red.find("def", name="zomg")
+    zomg.value = "return 42"
+    assert zomg.dumps() == """\
+def zomg():
+        return 42
+    """
+    assert len(zomg("endl")) == 2
+    assert zomg.value.node_list[-1].indent == "    "
 
 
 def test_set_attr_def_advanced_inline_dont_break_next_block_indent_one_endl():
     red = RedBaron(code_for_block_setattr)
-    red.find("def", name="zomg").value = "return 42\n"
-    assert len(red.find("def", name="zomg")("endl")) == 3
-    assert red.find("def", name="zomg").value.node_list[-1].indent == "    "
+    zomg = red.find("def", name="zomg")
+    zomg.value = "return 42\n"
+    assert zomg.dumps() == """\
+def zomg():
+        return 42
+    """
+    assert len(zomg("endl")) == 2
+    assert zomg.value.node_list[-1].indent == "    "
 
 
 def test_set_decorator_def():