round-trip sequence key elements0.12.14

1 files changed, 45 insertions, 2 deletions

diff --git a/comments.py b/comments.py
index 7dd8e8e..34154e8 100644
--- a/comments.py
+++ b/comments.py
@@ -12,7 +12,8 @@ from collections import MutableSet, Sized, Set  # type: ignore
 
 from ruamel.yaml.compat import ordereddict, PY2
 
-__all__ = ["CommentedSeq", "CommentedMap", "CommentedOrderedMap",
+__all__ = ["CommentedSeq", "CommentedKeySeq",
+           "CommentedMap", "CommentedOrderedMap",
            "CommentedSet", 'comment_attrib', 'merge_attrib']
 
 comment_attrib = '_yaml_comment'
@@ -321,8 +322,50 @@ class CommentedSeq(list, CommentedBase):
         return pre_comments
 
 
-class CommentedMapView(Sized):
+class CommentedKeySeq(tuple, CommentedBase):
+    """This primarily exists to be able to roundtrip keys that are sequences"""
+    def _yaml_add_comment(self, comment, key=NoComment):
+        if key is not NoComment:
+            self.yaml_key_comment_extend(key, comment)
+        else:
+            self.ca.comment = comment
+
+    def _yaml_add_eol_comment(self, comment, key):
+        self._yaml_add_comment(comment, key=key)
+
+    def _yaml_get_columnX(self, key):
+        return self.ca.items[key][0].start_mark.column
+
+    def _yaml_get_column(self, key):
+        column = None
+        sel_idx = None
+        pre, post = key-1, key+1
+        if pre in self.ca.items:
+            sel_idx = pre
+        elif post in self.ca.items:
+            sel_idx = post
+        else:
+            # self.ca.items is not ordered
+            for row_idx, k1 in enumerate(self):
+                if row_idx >= key:
+                    break
+                if row_idx not in self.ca.items:
+                    continue
+                sel_idx = row_idx
+        if sel_idx is not None:
+            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 CommentedMapView(Sized):
     __slots__ = '_mapping',
 
     def __init__(self, mapping):