diff --git a/redbaron/base_nodes.py b/redbaron/base_nodes.py index ff63540b..d8f227c8 100644 --- a/redbaron/base_nodes.py +++ b/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))]) diff --git a/redbaron/nodes.py b/redbaron/nodes.py index ffef00c6..c7cdcd99 100644 --- a/redbaron/nodes.py +++ b/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) diff --git a/tests/test_insert.py b/tests/test_insert.py index d3ac5697..c5d1b06c 100644 --- a/tests/test_insert.py +++ b/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): diff --git a/tests/test_setter.py b/tests/test_setter.py index e9e61296..33ec02e5 100644 --- a/tests/test_setter.py +++ b/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():