allow simple mappings as mapping keys

8 files changed, 172 insertions, 10 deletions

diff --git a/README.rst b/README.rst
index 2eeda47..c071fd0 100644
--- a/README.rst
+++ b/README.rst
@@ -54,8 +54,24 @@ ChangeLog
 
 .. should insert NEXT: at the beginning of line for next key (with empty line)
 
+NEXT:
+  - simple mappings can now be used as keys when round-tripping::
+
+      {a: 1, b: 2}: hello world
+      
+    although using the obvious operations (del, popitem) on the key will
+    fail, you can mutilate it by going through its attributes. If you load the
+    above YAML in `d`, then changing the value is cumbersome:
+
+        d = {CommentedKeyMap([('a', 1), ('b', 2)]): "goodbye"}
+    and changing the key even more so:
+
+        d[CommentedKeyMap([('b', 1), ('a', 2)])]= d.pop(CommentedKeyMap([('a', 1), ('b', 2)]))
+    (you can use a `dict` instead of a list of tuples (or ordereddict), but that might result
+    in a different order, of the keys of the key, in the output)
+
 0.15.57 (2018-08-15):
-  - Fix that CommentedSeq could no longer be used in adding or do a copy
+  - Fix that CommentedSeq could no longer be used in adding or do a sort
     (reported by `Christopher Wright <https://bitbucket.org/CJ-Wright4242/>`__)
 
 0.15.56 (2018-08-15):
diff --git a/_test/roundtrip.py b/_test/roundtrip.py
index 462dfc6..fa9d189 100644
--- a/_test/roundtrip.py
+++ b/_test/roundtrip.py
@@ -102,6 +102,7 @@ def round_trip(
     explicit_start=None,
     explicit_end=None,
     version=None,
+    dump_data=None,
 ):
     """
     inp:    input string to parse
@@ -113,6 +114,8 @@ def round_trip(
     if extra is not None:
         doutp += extra
     data = round_trip_load(inp, preserve_quotes=preserve_quotes)
+    if dump_data:
+        print('data', data)
     if intermediate is not None:
         if isinstance(intermediate, dict):
             for k, v in intermediate.items():
diff --git a/_test/test_fail.py b/_test/test_fail.py
index 2bb4230..2f90112 100644
--- a/_test/test_fail.py
+++ b/_test/test_fail.py
@@ -183,9 +183,51 @@ class TestFlowValues:
         """
         round_trip(inp)
 
-    # @pytest.mark.xfail(strict=True)
     def test_flow_value_with_colon_quoted(self):
         inp = """\
         {a: 'bcd:efg'}
         """
         round_trip(inp, preserve_quotes=True)
+
+
+class TestMappingKey:
+    def test_simple_mapping_key(self):
+        inp = """\
+        {a: 1, b: 2}: hello world
+        """
+        round_trip(inp, preserve_quotes=True, dump_data=False)
+
+    def test_set_simple_mapping_key(self):
+        from ruamel.yaml.comments import CommentedKeyMap
+
+        d = {CommentedKeyMap([('a', 1), ('b', 2)]): 'hello world'}
+        exp = dedent("""\
+        {a: 1, b: 2}: hello world
+        """)
+        assert round_trip_dump(d) == exp
+
+    def test_change_key_simple_mapping_key(self):
+        from ruamel.yaml.comments import CommentedKeyMap
+
+        inp = """\
+        {a: 1, b: 2}: hello world
+        """
+        d = round_trip_load(inp, preserve_quotes=True)
+        d[CommentedKeyMap([('b', 1), ('a', 2)])] = d.pop(CommentedKeyMap([('a', 1), ('b', 2)]))
+        exp = dedent("""\
+        {b: 1, a: 2}: hello world
+        """)
+        assert round_trip_dump(d) == exp
+
+    def test_change_value_simple_mapping_key(self):
+        from ruamel.yaml.comments import CommentedKeyMap
+
+        inp = """\
+        {a: 1, b: 2}: hello world
+        """
+        d = round_trip_load(inp, preserve_quotes=True)
+        d = {CommentedKeyMap([('a', 1), ('b', 2)]): 'goodbye'}
+        exp = dedent("""\
+        {a: 1, b: 2}: goodbye
+        """)
+        assert round_trip_dump(d) == exp
diff --git a/comments.py b/comments.py
index 09ba705..6713829 100644
--- a/comments.py
+++ b/comments.py
@@ -16,9 +16,9 @@ from ruamel.yaml.compat import ordereddict, PY2, string_types, MutableSliceableS
 from ruamel.yaml.scalarstring import ScalarString
 
 if PY2:
-    from collections import MutableSet, Sized, Set, MutableMapping
+    from collections import MutableSet, Sized, Set, MutableMapping, Mapping
 else:
-    from collections.abc import MutableSet, Sized, Set, MutableMapping
+    from collections.abc import MutableSet, Sized, Set, MutableMapping, Mapping
 
 if False:  # MYPY
     from typing import Any, Dict, Optional, List, Union, Optional  # NOQA
@@ -978,6 +978,94 @@ class CommentedMap(CommentedBase, MutableMapping):
         return res
 
 
+# based on brownie mappings
+@classmethod
+def raise_immutable(cls, *args, **kwargs):
+    raise TypeError('{} objects are immutable'.format(cls.__name__))
+
+
+class CommentedKeyMap(CommentedBase, Mapping):
+    __slots__ = Comment.attrib, '_od'
+    """This primarily exists to be able to roundtrip keys that are mappings"""
+
+    def __init__(self, *args, **kw):
+        # type: (Any, Any) -> None
+        if hasattr(self, '_od'):
+            raise_immutable(self)
+        self._od = ordereddict(*args, **kw)
+
+    __delitem__ = __setitem__ = clear = pop = popitem = setdefault = update = raise_immutable
+
+    # need to implement __getitem__, __iter__ and __len__
+    def __getitem__(self, index):
+        return self._od[index]
+
+    def __iter__(self):
+        for x in self._od.__iter__():
+            yield x
+
+    def __len__(self):
+        return len(self._od)
+
+    def __hash__(self):
+        return hash(tuple(self.items()))
+
+    def __repr__(self):
+        # type: () -> Any
+        if not hasattr(self, merge_attrib):
+            return self._od.__repr__()
+        return 'ordereddict(' + repr(list(self._items())) + ')'
+
+    @classmethod
+    def fromkeys(keys, v=None):
+        return CommentedKeyMap(dict.fromkeys(keys, v))
+
+    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)
+        else:
+            self.ca.comment = comment
+
+    def _yaml_add_eol_comment(self, comment, key):
+        # type: (Any, Any) -> None
+        self._yaml_add_comment(comment, key=key)
+
+    def _yaml_get_columnX(self, key):
+        # type: (Any) -> Any
+        return self.ca.items[key][0].start_mark.column
+
+    def _yaml_get_column(self, key):
+        # type: (Any) -> Any
+        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):
+        # type: () -> Any
+        pre_comments = []  # type: List[Any]
+        if self.ca.comment is None:
+            self.ca.comment = [None, pre_comments]
+        else:
+            self.ca.comment[1] = pre_comments
+        return pre_comments
+
+
 class CommentedOrderedMap(CommentedMap):
     __slots__ = (Comment.attrib,)
 
diff --git a/compat.py b/compat.py
index 5a39ca5..8b81975 100644
--- a/compat.py
+++ b/compat.py
@@ -92,7 +92,7 @@ if PY3:
     BytesIO = io.BytesIO
     # have unlimited precision
     no_limit_int = int
-    from collections.abc import Hashable, MutableSequence  # NOQA
+    from collections.abc import Hashable, MutableSequence, MutableMapping, Mapping  # NOQA
 
 else:
     string_types = basestring  # NOQA
@@ -111,7 +111,7 @@ else:
     BytesIO = cStringIO.StringIO
     # have unlimited precision
     no_limit_int = long  # NOQA not available on Python 3
-    from collections import Hashable, MutableSequence  # NOQA
+    from collections import Hashable, MutableSequence, MutableMapping, Mapping  # NOQA
 
 if False:  # MYPY
     # StreamType = Union[BinaryIO, IO[str], IO[unicode],  StringIO]
diff --git a/constructor.py b/constructor.py
index 8e1da70..9408c3c 100644
--- a/constructor.py
+++ b/constructor.py
@@ -17,10 +17,10 @@ from ruamel.yaml.nodes import *                               # NOQA
 from ruamel.yaml.nodes import (SequenceNode, MappingNode, ScalarNode)
 from ruamel.yaml.compat import (utf8, builtins_module, to_str, PY2, PY3,  # NOQA
                                 ordereddict, text_type, nprint, version_tnf, Hashable,
-                                MutableSequence)
+                                MutableSequence, MutableMapping)
 from ruamel.yaml.comments import *                               # NOQA
 from ruamel.yaml.comments import (CommentedMap, CommentedOrderedMap, CommentedSet,
-                                  CommentedKeySeq, CommentedSeq, TaggedScalar)
+                                  CommentedKeySeq, CommentedSeq, TaggedScalar, CommentedKeyMap)
 from ruamel.yaml.scalarstring import *                           # NOQA
 from ruamel.yaml.scalarstring import (PreservedScalarString, SingleQuotedScalarString,
                                       DoubleQuotedScalarString, ScalarString)
@@ -1330,7 +1330,13 @@ class RoundTripConstructor(SafeConstructor):
                     elif key_node.flow_style is False:
                         key_a.fa.set_block_style()
                     key = key_a
-                    # key = tuple(key)
+                elif isinstance(key, MutableMapping):
+                    key_a = CommentedKeyMap(key)
+                    if key_node.flow_style is True:
+                        key_a.fa.set_flow_style()
+                    elif key_node.flow_style is False:
+                        key_a.fa.set_block_style()
+                    key = key_a
             if PY2:
                 try:
                     hash(key)
diff --git a/emitter.py b/emitter.py
index 7fb572a..9f4da23 100644
--- a/emitter.py
+++ b/emitter.py
@@ -621,7 +621,9 @@ class Emitter(object):
                 self.write_pre_comment(self.event)
             self.write_indent()
             if self.check_simple_key():
-                if not isinstance(self.event, SequenceStartEvent):  # sequence keys
+                if not isinstance(
+                    self.event, (SequenceStartEvent, MappingStartEvent)
+                ):  # sequence keys
                     if self.event.style == '?':
                         self.write_indicator(u'?', True, indention=True)
                 self.states.append(self.expect_block_mapping_simple_value)
@@ -703,6 +705,7 @@ class Emitter(object):
         return length < self.MAX_SIMPLE_KEY_LENGTH and (
             isinstance(self.event, AliasEvent)
             or (isinstance(self.event, SequenceStartEvent) and self.event.flow_style is True)
+            or (isinstance(self.event, MappingStartEvent) and self.event.flow_style is True)
             or (
                 isinstance(self.event, ScalarEvent)
                 and not self.analysis.empty
diff --git a/representer.py b/representer.py
index c663397..9b97bdd 100644
--- a/representer.py
+++ b/representer.py
@@ -637,6 +637,7 @@ from ruamel.yaml.comments import (
     CommentedOrderedMap,
     CommentedSeq,
     CommentedKeySeq,
+    CommentedKeyMap,
     CommentedSet,
     comment_attrib,
     merge_attrib,
@@ -886,6 +887,9 @@ class RoundTripRepresenter(SafeRepresenter):
         if isinstance(data, CommentedKeySeq):
             self.alias_key = None
             return self.represent_sequence(u'tag:yaml.org,2002:seq', data, flow_style=True)
+        if isinstance(data, CommentedKeyMap):
+            self.alias_key = None
+            return self.represent_mapping(u'tag:yaml.org,2002:map', data, flow_style=True)
         return SafeRepresenter.represent_key(self, data)
 
     def represent_mapping(self, tag, mapping, flow_style=None):