preserve eol comments on flow style sequences/mappings

11 files changed, 111 insertions, 21 deletions

diff --git a/CHANGES b/CHANGES
index 25e4bf4..0449582 100644
--- a/CHANGES
+++ b/CHANGES
@@ -1,4 +1,10 @@
 
+0.7.2: 2015-03-29
+- support for end-of-line comments on flow style sequences and mappings
+
+0.7.1: 2015-03-27
+- RoundTrip capability of flow style sequences ( 'a: b, c, d' )
+
 0.7 2015-03-26
 - tests (currently failing) for inline sequece and non-standard spacing between
   block sequence dash and scalar (Anthony Sottile)
diff --git a/README.rst b/README.rst
index 5719f1b..2d34fe7 100644
--- a/README.rst
+++ b/README.rst
@@ -8,7 +8,8 @@ which supports YAML1.1
 
 Major differences with PyYAML 3.11:
 
-- intergrated Python 2 and 3 sources, running on Python 2.6, 2.7, 3.3 and 3.4.
+- intergrated Python 2 and 3 sources, running on Python 2.6, 2.7 (CPython, 
+  PyPy), 3.3 and 3.4.
 - round trip mode that **includes comments** (block mode, key ordering kept)
 - support for simple lists as mapping keys by transformation to tuples
 - ``!!omap`` generates ordereddict (C) on Python 2, collections.OrderedDict
diff --git a/py/__init__.py b/py/__init__.py
index 8ee5446..3757997 100644
--- a/py/__init__.py
+++ b/py/__init__.py
@@ -21,7 +21,7 @@ def _convert_version(tup):
     return ret_val
 
 
-version_info = (0, 7, 1)
+version_info = (0, 7, 2)
 __version__ = _convert_version(version_info)
 
 del _convert_version
diff --git a/py/composer.py b/py/composer.py
index b2c745d..981c5ef 100644
--- a/py/composer.py
+++ b/py/composer.py
@@ -113,13 +113,21 @@ class Composer(object):
                             start_event.start_mark, None,
                             flow_style=start_event.flow_style,
                             comment=start_event.comment)
+        #print('.................', node.comment)
         if anchor is not None:
             self.anchors[anchor] = node
         index = 0
+        #print('...............0.', node.comment)
         while not self.check_event(SequenceEndEvent):
+            #print('...............1.', node.comment)
             node.value.append(self.compose_node(node, index))
             index += 1
         end_event = self.get_event()
+        if node.flow_style is True and end_event.comment is not None:
+            if node.comment is not None:
+                print('Warning: unexpected end_event commment in sequence '
+                      'node {}'.format(node.flow_style))
+            node.comment = end_event.comment
         node.end_mark = end_event.end_mark
         self.check_end_doc_comment(end_event, node)
         return node
@@ -146,6 +154,8 @@ class Composer(object):
             # node.value[item_key] = item_value
             node.value.append((item_key, item_value))
         end_event = self.get_event()
+        if node.flow_style is True and end_event.comment is not None:
+            node.comment = end_event.comment
         node.end_mark = end_event.end_mark
         self.check_end_doc_comment(end_event, node)
         return node
diff --git a/py/emitter.py b/py/emitter.py
index 69d2c84..54edae1 100644
--- a/py/emitter.py
+++ b/py/emitter.py
@@ -258,7 +258,9 @@ class Emitter(object):
                 self.expect_scalar()
             elif isinstance(self.event, SequenceStartEvent):
                 if self.event.comment:
-                    self.write_post_comment(self.event)
+                    self.write_pre_comment(self.event)
+                    if self.event.flow_style is False and self.event.comment:
+                        self.write_post_comment(self.event)
                 # print('seq event', self.event)
                 if self.flow_level or self.canonical or self.event.flow_style \
                         or self.check_empty_sequence():
@@ -266,7 +268,7 @@ class Emitter(object):
                 else:
                     self.expect_block_sequence()
             elif isinstance(self.event, MappingStartEvent):
-                if self.event.comment:
+                if self.event.flow_style is False and self.event.comment:
                     self.write_post_comment(self.event)
                 if self.event.comment and self.event.comment[1]:
                     self.write_pre_comment(self.event)
@@ -318,6 +320,9 @@ class Emitter(object):
                 self.write_indicator(u',', False)
                 self.write_indent()
             self.write_indicator(u']', False)
+            if self.event.comment and self.event.comment[0]:
+                # eol comment on flow sequence
+                self.write_post_comment(self.event)
             self.state = self.states.pop()
         else:
             self.write_indicator(u',', False)
@@ -339,6 +344,9 @@ class Emitter(object):
             self.indent = self.indents.pop()
             self.flow_level -= 1
             self.write_indicator(u'}', False)
+            #if self.event.comment and self.event.comment[0]:
+            #    # eol comment on flow sequence
+            #    self.write_post_comment(self.event)
             self.state = self.states.pop()
         else:
             if self.canonical or self.column > self.best_width:
@@ -361,6 +369,9 @@ class Emitter(object):
                 self.write_indicator(u',', False)
                 self.write_indent()
             self.write_indicator(u'}', False)
+            if self.event.comment and self.event.comment[0]:
+                # eol comment on flow mapping
+                self.write_post_comment(self.event)
             self.state = self.states.pop()
         else:
             self.write_indicator(u',', False)
@@ -1201,6 +1212,8 @@ class Emitter(object):
 
     def write_pre_comment(self, event):
         comments = event.comment[1]
+        if comments is None:
+            return
         try:
             for comment in comments:
                 if isinstance(event, MappingStartEvent) and \
diff --git a/py/parser.py b/py/parser.py
index 71258d0..bb90f2a 100644
--- a/py/parser.py
+++ b/py/parser.py
@@ -375,10 +375,16 @@ class Parser(object):
                     end_mark = self.peek_token().start_mark
                     # should inserting the comment be dependent on the
                     # indentation?
-                    comment = self.peek_token().split_comment()
+                    pt = self.peek_token()
+                    comment = pt.comment
+                    #print('pt0', type(pt))
+                    if comment is None or comment[1] is None:
+                        comment = pt.split_comment()
+                    #print('pt1', comment)
                     event = SequenceStartEvent(
                         anchor, tag, implicit, start_mark, end_mark,
-                        flow_style=False, comment=comment,
+                        flow_style=False,
+                        comment=comment,
                     )
                     self.state = self.parse_block_sequence_first_entry
                 elif block and self.check_token(BlockMappingStartToken):
@@ -412,7 +418,7 @@ class Parser(object):
     def parse_block_sequence_first_entry(self):
         token = self.get_token()
         # move any comment from start token
-        token.move_comment(self.peek_token())
+        #token.move_comment(self.peek_token())
         self.marks.append(token.start_mark)
         return self.parse_block_sequence_entry()
 
@@ -420,6 +426,8 @@ class Parser(object):
         if self.check_token(BlockEntryToken):
             token = self.get_token()
             token.move_comment(self.peek_token())
+            #print('================== here be dragons', repr(token))
+            #print(token.comment)
             if not self.check_token(BlockEntryToken, BlockEndToken):
                 self.states.append(self.parse_block_sequence_entry)
                 return self.parse_block_node()
@@ -622,7 +630,8 @@ class Parser(object):
                 self.states.append(self.parse_flow_mapping_empty_value)
                 return self.parse_flow_node()
         token = self.get_token()
-        event = MappingEndEvent(token.start_mark, token.end_mark)
+        event = MappingEndEvent(token.start_mark, token.end_mark,
+                                comment=token.comment)
         self.state = self.states.pop()
         self.marks.pop()
         return event
diff --git a/py/scanner.py b/py/scanner.py
index faa161a..721437c 100644
--- a/py/scanner.py
+++ b/py/scanner.py
@@ -1520,6 +1520,7 @@ class RoundTripScanner(Scanner):
             return self.tokens[0]
 
     def _gather_comments(self):
+        """combine multiple comment lines"""
         comments = []
         if not self.tokens:
             return comments
@@ -1552,11 +1553,17 @@ class RoundTripScanner(Scanner):
             self.fetch_more_tokens()
         self._gather_comments()
         if self.tokens:
-            # only add post comment to scalar token or value token. otherwise
+            # only add post comment to single line tokens:
+            # scalar, value token. FlowXEndToken, otherwise
             # hidden streamtokens could get them (leave them and they will be
             # pre comments for the next map/seq
             if len(self.tokens) > 1 and \
-               isinstance(self.tokens[0], (ScalarToken, ValueToken)) and \
+               isinstance(self.tokens[0], (
+                   ScalarToken,
+                   ValueToken,
+                   FlowSequenceEndToken,
+                   FlowMappingEndToken,
+                   )) and \
                isinstance(self.tokens[1], CommentToken) and \
                self.tokens[0].end_mark.line == self.tokens[1].start_mark.line:
                 self.tokens_taken += 1
diff --git a/py/serializer.py b/py/serializer.py
index 38db4c5..7e891e1 100644
--- a/py/serializer.py
+++ b/py/serializer.py
@@ -104,7 +104,13 @@ class Serializer(object):
                 implicit = (node.tag
                             == self.resolve(SequenceNode, node.value, True))
                 comment = node.comment
+                # print('comment >>>>>>>>>>>>>.', comment, node.flow_style)
                 end_comment = None
+                seq_comment = None
+                if node.flow_style is True:
+                    if comment:  # eol comment on flow style sequence
+                        seq_comment = comment[0]
+                        # comment[0] = None
                 if comment and len(comment) > 2:
                     end_comment = comment[2]
                 else:
@@ -116,12 +122,17 @@ class Serializer(object):
                 for item in node.value:
                     self.serialize_node(item, node, index)
                     index += 1
-                self.emit(SequenceEndEvent(comment=[None, end_comment]))
+                self.emit(SequenceEndEvent(comment=[seq_comment, end_comment]))
             elif isinstance(node, MappingNode):
                 implicit = (node.tag
                             == self.resolve(MappingNode, node.value, True))
                 comment = node.comment
                 end_comment = None
+                map_comment = None
+                if node.flow_style is True:
+                    if comment:  # eol comment on flow style sequence
+                        map_comment = comment[0]
+                        # comment[0] = None
                 if comment and len(comment) > 2:
                     end_comment = comment[2]
                 self.emit(MappingStartEvent(alias, node.tag, implicit,
@@ -130,5 +141,5 @@ class Serializer(object):
                 for key, value in node.value:
                     self.serialize_node(key, node, None)
                     self.serialize_node(value, node, key)
-                self.emit(MappingEndEvent(comment=[None, end_comment]))
+                self.emit(MappingEndEvent(comment=[map_comment, end_comment]))
             self.ascend_resolver()
diff --git a/py/yaml.py b/py/yaml.py
index e92f135..edfb338 100644
--- a/py/yaml.py
+++ b/py/yaml.py
@@ -98,7 +98,7 @@ class YAML:
             print('-' * 15)
 
         def print_tokens(input):
-            print('Tokens ' + '#' * 60)
+            print('Tokens (from scanner) ' + '#' * 50)
             tokens = ruamel.yaml.scan(input, ruamel.yaml.RoundTripLoader)
             for idx, token in enumerate(tokens):
                 # print(token.start_mark)
@@ -116,13 +116,13 @@ class YAML:
             print(ruamel.yaml.serialize(nodes, indent=False, Dumper=dumper))
 
         def print_events(input):
-            print('Events ' + '#' * 60)
+            print('Events (from parser) ' + '#' * 50)
             events = ruamel.yaml.parse(input, ruamel.yaml.RoundTripLoader)
             for idx, event in enumerate(events):
                 print("{0:2} {1}".format(idx, event))
 
         def print_nodes(input):
-            print('Nodes ' + '#' * 60)
+            print('Nodes (from composer) ' + '#' * 50)
             x = ruamel.yaml.compose(input, ruamel.yaml.RoundTripLoader)
             x.dump()  # dump the node
 
@@ -142,7 +142,33 @@ class YAML:
           secure: optional
         """)
 
-        input = "a: [b, c, d]"
+        input = dedent("""\
+        # comment A
+        a: [b, c, d]  # comment B
+        j: [k, l, m]  # comment C
+        """)
+
+        Xinput = dedent("""\
+        # comment A
+        - {a: 1, b: hallo}  # comment B
+        - {j: fka, k: 42}   # comment C
+        """)
+
+        Xinput = dedent("""\
+            # comment A
+            - {a: 1, b: hallo}
+            - {j: fka, k: 42}
+            """)
+
+        Xinput = dedent("""\
+        # C start a
+        # C start b
+        - abc      # abc comment
+        - ghi
+        - klm      # klm comment
+        # C end a
+        # C end b
+        """)
 
         print_input(input)
         print_tokens(input)
diff --git a/test/test_indentation.py b/test/test_indentation.py
index 67e2951..9a79c64 100644
--- a/test/test_indentation.py
+++ b/test/test_indentation.py
@@ -40,10 +40,7 @@ def test_roundtrip_mapping_of_inline_lists_comments():
     output = rt(s)
     assert s == output
 
-# the following doesn't work correctly. The comment for the sequence is
-# emitted on the next line instead of after the flow sequence
-@pytest.mark.xfail
-def test_roundtrip_mapping_of_inline_lists_eol_comments():
+def test_roundtrip_mapping_of_inline_sequence_eol_comments():
     s = dedent("""\
     # comment A
     a: [a, b, c]  # comment B
@@ -80,6 +77,16 @@ def test_roundtrip_flow_mapping():
     output = ruamel.yaml.dump(data, Dumper=ruamel.yaml.RoundTripDumper)
     assert s == output
 
+def test_roundtrip_sequence_of_inline_mappings_eol_comments():
+    s = dedent("""\
+    # comment A
+    - {a: 1, b: hallo}  # comment B
+    - {j: fka, k: 42}  # comment C
+    """)
+    output = rt(s)
+    assert s == output
+
+
 
 ############# indentation
 
diff --git a/tox.ini b/tox.ini
index 414a1c9..02c9109 100644
--- a/tox.ini
+++ b/tox.ini
@@ -3,7 +3,7 @@ envlist = py27,py34,py26,py33,pypy
 
 [testenv]
 commands =
-    py.test -x test
+    py.test test
 deps =
     pytest
     configobj