Merge pull request #35 from tompaton/issue_32

Add test+fix for issue #32 and xfail-marked test for #33
pycontribs · Feb 8, 2021 · 12fbe3d · 12fbe3d
2 parents 79a520b + 666075b
commit 12fbe3d
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 1 deletion.
diff --git a/_test/test_comment_manipulation.py b/_test/test_comment_manipulation.py
@@ -637,3 +637,87 @@ def test_map_set_comment_before_and_after_non_first_key_02(self):
             test3: 3
         """
         compare(data, exp)
+
+    # issue 32
+    def test_yaml_add_eol_comment_issue_32(self):
+        data = load(
+            """
+        items:
+        - one: 1
+          uno: '1'
+        -  # item 2
+          two: 2
+          duo: '2'
+        - three: 3
+        """
+        )
+
+        data['items'].yaml_add_eol_comment('second pass', key=1)
+
+        exp = """
+        items:
+        - one: 1
+          uno: '1'
+        -  # second pass
+          two: 2
+          duo: '2'
+        - three: 3
+        """
+
+        compare(data, exp)
+
+    def test_yaml_add_eol_comment_issue_32_ok(self):
+        data = load(
+            """
+        items:
+        - one
+        - two  # item 2
+        - three
+        """
+        )
+
+        data['items'].yaml_add_eol_comment('second pass', key=1)
+
+        exp = """
+        items:
+        - one
+        - two  # second pass
+        - three
+        """
+
+        compare(data, exp)
+
+    # issue 33
+    @pytest.mark.xfail(reason="open issue", raises=AttributeError)
+    def test_yaml_set_start_comment_issue_33(self):
+        data = load(
+            """
+        items:
+        # item 1
+        - one: 1
+          uno: '1'
+        # item 2
+        - two: 2
+          duo: '2'
+        # item 3
+        - three: 3
+        """
+        )
+
+        data['items'][0].yaml_set_start_comment('uno')
+        data['items'][1].yaml_set_start_comment('duo')
+        data['items'][2].yaml_set_start_comment('tre')
+
+        exp = """
+        items:
+        # uno
+        - one: 1
+          uno: '1'
+        # duo
+        - two: 2
+          duo: '2'
+        # tre
+        - three: 3
+        """
+
+        compare(data, exp)
diff --git a/lib/ruyaml/comments.py b/lib/ruyaml/comments.py
@@ -426,7 +426,7 @@ def __eq__(self, other):
     def _yaml_add_comment(self, comment, key=NoComment):
         # type: (Any, Optional[Any]) -> None
         if key is not NoComment:
-            self.yaml_key_comment_extend(key, comment)
+            self.yaml_key_comment_extend(key, comment, clear=True)
         else:
             self.ca.comment = comment