- set anchor explicitly (yaml_set_anchor)

- add before item comment (yaml_set_start_comment)

6 files changed, 228 insertions, 10 deletions

diff --git a/py/__init__.py b/py/__init__.py
index ebe529f..1c14001 100644
--- a/py/__init__.py
+++ b/py/__init__.py
@@ -21,7 +21,7 @@ def _convert_version(tup):
     return ret_val
 
 
-version_info = (0, 10)
+version_info = (0, 10, 1)
 __version__ = _convert_version(version_info)
 
 del _convert_version
diff --git a/py/comments.py b/py/comments.py
index 577b984..0db9457 100644
--- a/py/comments.py
+++ b/py/comments.py
@@ -104,6 +104,7 @@ class Anchor(object):
 
     def __init__(self):
         self.value = None
+        self.always_dump = False
 
 class CommentedBase(object):
     @property
@@ -137,6 +138,19 @@ class CommentedBase(object):
             l[3].extend(comment[0])
         l[2] = comment[0]
 
+    def yaml_set_start_comment(self, comment, indent=0):
+        """overwrites any preceding comment lines on an object
+        expects comment to be without `#` and possible have mutlple lines
+        """
+        from .error import Mark
+        from .tokens import CommentToken
+        pre_comments = self._yaml_get_pre_comment()
+        if comment[-1] == '\n':
+            comment = comment[:-1]  # strip final newline if there
+        start_mark = Mark(None, None, None, indent, None, None)
+        for com in comment.split('\n'):
+            pre_comments.append(CommentToken('# ' + com + '\n', start_mark, None))
+
     @property
     def fa(self):
         if not hasattr(self, Format.attrib):
@@ -182,12 +196,11 @@ class CommentedBase(object):
     def yaml_anchor(self):
         if not hasattr(self, Anchor.attrib):
             return None
-        return self.anchor.value
+        return self.anchor
 
-    def set_yaml_anchor(self, value):
+    def yaml_set_anchor(self, value, always_dump=False):
         self.anchor.value = value
-
-
+        self.anchor.always_dump = always_dump
 
 class CommentedSeq(list, CommentedBase):
     __slots__ = [Comment.attrib, ]
@@ -224,6 +237,14 @@ class CommentedSeq(list, CommentedBase):
             column = self._yaml_get_columnX(sel_idx)
         return column
 
+    def _yaml_get_pre_comment(self):
+        if self.ca.comment is None:
+            pre_comments = []
+            self.ca.comment = [None, pre_comments]
+        else:
+            pre_comments = self.ca.comment[1] = []
+        return pre_comments
+
 
 class CommentedMap(ordereddict, CommentedBase):
     __slots__ = [Comment.attrib, ]
@@ -272,6 +293,14 @@ class CommentedMap(ordereddict, CommentedBase):
             column = self._yaml_get_columnX(sel_idx)
         return column
 
+    def _yaml_get_pre_comment(self):
+        if self.ca.comment is None:
+            pre_comments = []
+            self.ca.comment = [None, pre_comments]
+        else:
+            pre_comments = self.ca.comment[1] = []
+        return pre_comments
+
     def update(self, vals):
         try:
             ordereddict.update(self, vals)
diff --git a/py/constructor.py b/py/constructor.py
index 274aa22..bfcd17d 100644
--- a/py/constructor.py
+++ b/py/constructor.py
@@ -911,7 +911,7 @@ class RoundTripConstructor(SafeConstructor):
         if node.anchor:
             from ruamel.yaml.serializer import templated_id
             if not templated_id(node.anchor):
-                maptyp.set_yaml_anchor(node.anchor)
+                maptyp.yaml_set_anchor(node.anchor)
         for key_node, value_node in node.value:
             # keys can be list -> deep
             key = self.construct_object(key_node, deep=True)
diff --git a/py/serializer.py b/py/serializer.py
index 30c80ba..0bdc558 100644
--- a/py/serializer.py
+++ b/py/serializer.py
@@ -77,7 +77,13 @@ class Serializer(object):
             if self.anchors[node] is None:
                 self.anchors[node] = self.generate_anchor(node)
         else:
-            self.anchors[node] = None
+            anchor = None
+            try:
+                if node.anchor.always_dump:
+                    anchor = node.anchor.value
+            except:
+                pass
+            self.anchors[node] = anchor
             if isinstance(node, SequenceNode):
                 for item in node.value:
                     self.anchor_node(item)
@@ -88,7 +94,7 @@ class Serializer(object):
 
     def generate_anchor(self, node):
         try:
-            anchor = node.anchor
+            anchor = node.anchor.value
         except:
             anchor = None
         if anchor is None:
diff --git a/test/test_anchor.py b/test/test_anchor.py
index 477c8ee..422b9d6 100644
--- a/test/test_anchor.py
+++ b/test/test_anchor.py
@@ -69,7 +69,8 @@ class TestAnchorsAliases:
         assert d.yaml_anchor() is None  # got dropped as it matches pattern
         e = data['e']
         assert isinstance(e, CommentedMap)
-        assert e.yaml_anchor() == 'etemplate'
+        assert e.yaml_anchor().value == 'etemplate'
+        assert e.yaml_anchor().always_dump is False
 
     #@pytest.mark.xfail
     def test_anchor_id_retained(self):
@@ -156,4 +157,20 @@ class TestAnchorsAliases:
 
     def test_merge_01(self):
         data = load(self.merge_yaml)
-        compare(data, self.merge_yaml)
\ No newline at end of file
+        compare(data, self.merge_yaml)
+
+    def test_add_anchor(self):
+        from ruamel.yaml.comments import CommentedMap
+        data = CommentedMap()
+        data_a = CommentedMap()
+        data['a'] = data_a
+        data_a['c'] = 3
+        data['b'] = 2
+        data.yaml_set_anchor('klm', always_dump=True)
+        data['a'].yaml_set_anchor('xyz', always_dump=True)
+        compare(data, """
+        &klm
+        a: &xyz
+          c: 3
+        b: 2
+        """)
diff --git a/test/test_comment_manipulation.py b/test/test_comment_manipulation.py
index 128cade..a500b79 100644
--- a/test/test_comment_manipulation.py
+++ b/test/test_comment_manipulation.py
@@ -1,5 +1,7 @@
 # coding: utf-8
 
+from __future__ import print_function
+
 import pytest
 from textwrap import dedent
 
@@ -225,3 +227,167 @@ class TestCommentsManipulation:
             e: 5     # comment 3
             """)
 
+# the ugly {comment} in the following is because
+# py.test cannot handle comments in strings properly
+# https://bitbucket.org/pytest-dev/pytest/issue/752/internalerror-indexerror-list-index-out-of
+
+    #@pytest.mark.xfail
+    def test_before_top_map_rt(self):
+        data = load("""
+        a: 1
+        b: 2
+        """)
+        data.yaml_set_start_comment('Hello\nWorld\n')
+        compare(data, """
+        {comment} Hello
+        {comment} World
+        a: 1
+        b: 2
+        """.format(comment='#'))
+
+    def test_before_top_map_replace(self):
+        data = load("""
+        # abc
+        # def
+        a: 1 # 1
+        b: 2
+        """)
+        data.yaml_set_start_comment('Hello\nWorld\n')
+        compare(data, """
+        {comment} Hello
+        {comment} World
+        a: 1 # 1
+        b: 2
+        """.format(comment='#'))
+
+    def test_before_top_map_from_scratch(self):
+        from ruamel.yaml.comments import CommentedMap
+        data = CommentedMap()
+        data['a'] = 1
+        data['b'] = 2
+        data.yaml_set_start_comment('Hello\nWorld\n')
+        #print(data.ca)
+        #print(data.ca._items)
+        compare(data, """
+            {comment} Hello
+            {comment} World
+            a: 1
+            b: 2
+            """.format(comment='#'))
+
+    def test_before_top_seq_rt(self):
+        data = load("""
+        - a
+        - b
+        """)
+        data.yaml_set_start_comment('Hello\nWorld\n')
+        print(round_trip_dump(data))
+        compare(data, """
+        # Hello
+        # World
+        - a
+        - b
+        """)
+
+    def test_before_top_seq_rt_replace(self):
+        data = load("""
+        {comment} this
+        {comment} that
+        - a
+        - b
+        """.format(comment='#'))
+        data.yaml_set_start_comment('Hello\nWorld\n')
+        print(round_trip_dump(data))
+        compare(data, """
+        {comment} Hello
+        {comment} World
+        - a
+        - b
+        """.format(comment='#'))
+
+    def test_before_top_seq_from_scratch(self):
+        from ruamel.yaml.comments import CommentedSeq
+        data = CommentedSeq()
+        data.append('a')
+        data.append('b')
+        data.yaml_set_start_comment('Hello\nWorld\n')
+        print(round_trip_dump(data))
+        compare(data, """
+        {comment} Hello
+        {comment} World
+        - a
+        - b
+        """.format(comment='#'))
+
+    # nested variants
+    def test_before_nested_map_rt(self):
+        data = load("""
+        a: 1
+        b:
+          c: 2
+          d: 3
+        """)
+        data['b'].yaml_set_start_comment('Hello\nWorld\n')
+        compare(data, """
+        a: 1
+        b:
+        {comment} Hello
+        {comment} World
+          c: 2
+          d: 3
+        """.format(comment='#'))
+
+    def test_before_nested_map_rt_indent(self):
+        data = load("""
+        a: 1
+        b:
+          c: 2
+          d: 3
+        """)
+        data['b'].yaml_set_start_comment('Hello\nWorld\n', indent=2)
+        compare(data, """
+        a: 1
+        b:
+          {comment} Hello
+          {comment} World
+          c: 2
+          d: 3
+        """.format(comment='#'))
+        print(data['b'].ca)
+
+    def test_before_nested_map_from_scratch(self):
+        from ruamel.yaml.comments import CommentedMap
+        data = CommentedMap()
+        datab = CommentedMap()
+        data['a'] = 1
+        data['b'] = datab
+        datab['c'] = 2
+        datab['d'] = 3
+        data['b'].yaml_set_start_comment('Hello\nWorld\n')
+        compare(data, """
+        a: 1
+        b:
+        {comment} Hello
+        {comment} World
+          c: 2
+          d: 3
+        """.format(comment='#'))
+
+    def test_before_nested_seq_from_scratch(self):
+        from ruamel.yaml.comments import CommentedMap, CommentedSeq
+        data = CommentedMap()
+        datab = CommentedSeq()
+        data['a'] = 1
+        data['b'] = datab
+        datab.append('c')
+        datab.append('d')
+        data['b'].yaml_set_start_comment('Hello\nWorld\n', indent=2)
+        compare(data, """
+        a: 1
+        b:
+          {comment} Hello
+          {comment} World
+          - c
+          - d
+        """.format(comment='#'))
+