fix issues with comments

- tailing key comments and leading value comments were lost if
  rest-of-line strings were used
- inline comments in multiline keys were lost

Author: Ken Kundert

doc/comments.rst

@@ -120,8 +120,6 @@ comment for the first data item.  The partition is at the *last blank
 line* in the buffer.  If there is no blank line, the entire content is
 leading on the first data item (no header).
 
-Comments in a document that contains no data are all header comments.
-
 Leading / trailing comments
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -416,6 +414,10 @@ the body and the footer comments.  See :meth:`Location.set_spacing` for
 how to attach a *spacing* dict to a particular Location, replacing the
 global spacing within that subtree.
 
+You can also set the spacing on :class:`Location` objects in the keymap 
+directly, which allows you to specify different spacing rules for different 
+parts of the document.
+
 When the load and dump happen in different processes (or are otherwise
 separated in time), use :func:`keymap_to_jsonable` and
 :func:`keymap_from_jsonable` to ship the keymap between them as plain

nestedtext/nestedtext.py

@@ -1472,18 +1472,57 @@ def _add_keymap(self, keys, location):
             # each Location own its comments outright -- two Locations that
             # share a Line (parent and its first child) cannot then
             # double-emit the same comment block.
-            kl = location.key_line if location.key_line is not None else location.line
+            key_first = (
+                location.key_line if location.key_line is not None
+                else location.line
+            )
             vl = location.line
             ve = location.value_end_line
-            if kl is not None and kl.leading_comments:
-                location.key_leading_comments = list(kl.leading_comments)
-                kl.leading_comments = []
-            if vl is not None and vl is not kl and vl.leading_comments:
+            # For a multi-line key, walk the chain of fragment lines via
+            # next_line (set by Lines.read_lines only between consecutive
+            # ``key item`` lines at the same depth) so that comments
+            # staged on later fragments are claimed too.
+            key_last = key_first
+            if key_first is not None:
+                while (
+                    getattr(key_last, "next_line", None) is not None
+                    and key_last.next_line.kind == "key item"
+                    and key_last.next_line.depth == key_last.depth
+                ):
+                    key_last = key_last.next_line
+            if key_first is not None and key_first.leading_comments:
+                location.key_leading_comments = list(key_first.leading_comments)
+                key_first.leading_comments = []
+            if (
+                vl is not None
+                and vl is not key_first
+                and vl is not key_last
+                and vl.leading_comments
+            ):
                 location.value_leading_comments = list(vl.leading_comments)
                 vl.leading_comments = []
-            if kl is not None and kl is not ve and kl.trailing_comments:
-                location.key_trailing_comments = list(kl.trailing_comments)
-                kl.trailing_comments = []
+            # key_trailing collects (a) leading_comments staged on each
+            # *intermediate* key-fragment line -- these are comments that
+            # appeared between fragments of the multi-line key, the
+            # multi-line-key analogue of inline-in-multi-line-string -- and
+            # (b) the trailing_comments on each fragment, including the
+            # last.  All are emitted at the key-trailing position.  When
+            # the entire key+value is on one line (key_first == ve), the
+            # trailing comments belong to value_trailing instead.
+            kt = []
+            cur = key_first
+            while cur is not None:
+                if cur is not key_first and cur.leading_comments:
+                    kt.extend(cur.leading_comments)
+                    cur.leading_comments = []
+                if cur is not ve and cur.trailing_comments:
+                    kt.extend(cur.trailing_comments)
+                    cur.trailing_comments = []
+                if cur is key_last:
+                    break
+                cur = cur.next_line
+            if kt:
+                location.key_trailing_comments = kt
             if ve is not None and ve.trailing_comments:
                 location.value_trailing_comments = list(ve.trailing_comments)
                 ve.trailing_comments = []
@@ -2079,6 +2118,16 @@ def render_dict_item(self, key, value, keys, values):
             or key[:2] in ["- ", "> ", ": "]
             or ": " in key
         )
+        # The key_trailing and value_leading comment slots only have a
+        # rendering position in the multi-line dict-item *value* form
+        # (between the key line and the value's first line).  If either
+        # slot has any contribution -- static or via a parent provider --
+        # force the value onto its own line so those comments don't get
+        # silently dropped.
+        force_multiline_value = (
+            not multiline_key_required
+            and self._comments_force_multiline(keys)
+        )
         if multiline_key_required:
             key = "\n".join(": "+l if l else ":" for l in key.split("\n"))
             if self.is_a_dict(value) or self.is_a_list(value):
@@ -2089,8 +2138,39 @@ def render_dict_item(self, key, value, keys, values):
             else:
                 value = self.render_value(value, keys, values)
             return key + "\n" + add_leader(value, self.indent*" " + "> ")
-        else:
-            return add_prefix(key + ":", self.render_value(value, keys, values))
+        if force_multiline_value:
+            # Plain "key:" syntax, but force the value onto its own line
+            # so key_trailing / value_leading have a place to render.
+            if self.is_a_dict(value) or self.is_a_list(value):
+                return key + ":" + self.render_value(value, keys, values)
+            if is_str(value):
+                value_text = convert_line_terminators(value)
+            else:
+                value_text = self.render_value(value, keys, values)
+            return key + ":\n" + add_leader(value_text, self.indent*" " + "> ")
+        return add_prefix(key + ":", self.render_value(value, keys, values))
+
+    # _comments_force_multiline {{{3
+    def _comments_force_multiline(self, keys):
+        """Return True if any source -- static key_trailing/value_leading
+        on this Location, or a parent provider for either slot -- will
+        contribute Comments that need the multi-line dict-item form.
+        """
+        if not is_mapping(self.map_keys):
+            return False
+        loc = self.map_keys.get(keys)
+        if loc is not None:
+            if loc.get_key_trailing_comments() or loc.get_value_leading_comments():
+                return True
+        if keys:
+            parent_loc = self.map_keys.get(keys[:-1])
+            if parent_loc is not None:
+                if (
+                    parent_loc.get_key_trailing_provider() is not None
+                    or parent_loc.get_value_leading_provider() is not None
+                ):
+                    return True
+        return False
 
     # render_inline_value {{{3
     def render_inline_value(self, obj, exclude, keys, values):
@@ -2318,15 +2398,36 @@ def _wrap_with_comments(self, rendered_value, keys):
         value_leading = self._comments_to_lines(vl, natural=val_natural)
         trailing = self._comments_to_lines(vt, natural=val_natural)
         value_lines = rendered_value.split("\n")
-        # Inject key_trailing and value_leading between the key line (first
-        # line) and the value's first line.  Only meaningful when the item
-        # spans multiple lines.
+        # Inject key_trailing and value_leading between the rendered key
+        # (which may span several lines for multi-line keys) and the
+        # value's first line.  We detect the key's line count by counting
+        # consecutive leading lines that look like multi-line key
+        # fragments (``: frag`` or ``:`` after lstrip) at the *same*
+        # indent.  If no multi-line-key prefix is present, the key is the
+        # first line (e.g. ``key:``).  Inline values (single-line output)
+        # don't get key_trailing / value_leading -- those are forced into
+        # multi-line by ``_comments_force_multiline``, which is what
+        # ensures we never silently drop them here.
         if (key_trailing or value_leading) and len(value_lines) > 1:
+            boundary = 0
+            key_indent = None
+            for line in value_lines:
+                stripped = line.lstrip()
+                if not (stripped == ":" or stripped.startswith(": ")):
+                    break
+                indent = len(line) - len(stripped)
+                if key_indent is None:
+                    key_indent = indent
+                elif indent != key_indent:
+                    break
+                boundary += 1
+            if boundary == 0:
+                boundary = 1
             value_lines = (
-                value_lines[:1]
+                value_lines[:boundary]
                 + key_trailing
                 + value_leading
-                + value_lines[1:]
+                + value_lines[boundary:]
             )
         return "\n".join(leading + value_lines + trailing)
 

tests/test_comments.py

@@ -1096,6 +1096,173 @@ def header(k):
     assert "# start of 2024-02" not in out
 
 
+def test_multi_line_key_round_trip_preserves_all_comment_slots():
+    """A multi-line key with key_trailing, value_leading, and value_trailing
+    must round-trip without dropping or mispositioning any of them."""
+    src = (
+        "# heading\n"
+        "\n"
+        "# leading on key\n"
+        ": key1a\n"
+        ": key1b\n"
+        "        # trailing on key\n"
+        "    # leading on value\n"
+        "    > nutz\n"
+        "        # trailing on value\n"
+        "\n"
+        "# footer\n"
+    )
+    keymap = {}
+    nt.loads(src, top="dict", keymap=keymap)
+    # The multi-line key is "key1a\nkey1b".
+    loc = keymap[("key1a\nkey1b",)]
+    assert [c.text for c in loc.get_key_leading_comments()]    == ["leading on key"]
+    assert [c.text for c in loc.get_key_trailing_comments()]   == ["trailing on key"]
+    assert [c.text for c in loc.get_value_leading_comments()]  == ["leading on value"]
+    assert [c.text for c in loc.get_value_trailing_comments()] == ["trailing on value"]
+
+
+def test_multi_line_key_inline_comment_collected_as_key_trailing():
+    """A comment between fragments of a multi-line key is collected and
+    emitted at the key_trailing position (similar to the inline-in-
+    multi-line-string convention)."""
+    src = (
+        ": key1a\n"
+        "# inline in key\n"
+        ": key1b\n"
+        "    > value\n"
+    )
+    keymap = {}
+    nt.loads(src, top="dict", keymap=keymap)
+    loc = keymap[("key1a\nkey1b",)]
+    # the inline comment ends up in key_trailing
+    texts = [c.text for c in loc.get_key_trailing_comments()]
+    assert "inline in key" in texts
+
+
+def test_dumper_multi_line_key_with_inner_multi_line_key():
+    """The dumper's key-boundary scan must stop when it encounters a
+    deeper-indented ':' line -- that belongs to the inner dict's
+    multi-line key, not to the outer's."""
+    keymap = {}
+    annotate(("a\nb",), keymap, key_trailing=[Comment("outer kt", tab=1)])
+    data = {"a\nb": {"c\nd": "v"}}
+    out = nt.dumps(data, map_keys=keymap, indent=4)
+    lines = out.split("\n")
+    outer_a = lines.index(": a")
+    outer_b = lines.index(": b")
+    kt = lines.index("        # outer kt")
+    inner_c = lines.index("    : c")
+    inner_d = lines.index("    : d")
+    # outer kt goes between the outer key's last fragment and the inner
+    # dict's first fragment.
+    assert outer_a < outer_b < kt < inner_c < inner_d
+
+
+def test_dumper_multi_line_key_places_kt_vl_after_all_fragments():
+    """When emitting a multi-line key, key_trailing/value_leading
+    comments go AFTER the last key fragment, not between the
+    fragments."""
+    keymap = {}
+    annotate(("a\nb",), keymap,
+        key_trailing=[Comment("kt", tab=1)],
+        value_leading=[Comment("vl", tab=0)],
+    )
+    data = {"a\nb": "v"}
+    out = nt.dumps(data, map_keys=keymap, indent=4)
+    lines = out.split("\n")
+    a = lines.index(": a")
+    b = lines.index(": b")
+    kt = lines.index("        # kt")
+    vl = lines.index("    # vl")
+    val = lines.index("    > v")
+    # all key fragments precede the comments; comments precede the value
+    assert a < b < kt < vl < val
+
+
+def test_round_trip_preserves_key_trailing_and_value_leading():
+    """A load with key_trailing and value_leading comments on a scalar
+    value must round-trip exactly -- the dumper must NOT collapse to
+    inline form (which would drop those slots)."""
+    src = (
+        "# heading comment for document\n"
+        "\n"
+        "# leading comment for key\n"
+        "key:\n"
+        "        # trailing comment for key\n"
+        "    # leading comment for value\n"
+        "    > nutz\n"
+        "        # trailing comment for value\n"
+        "\n"
+        "# footer comment for document\n"
+    )
+    keymap = {}
+    data = nt.loads(src, top="any", keymap=keymap)
+    out = nt.dumps(data, map_keys=keymap)
+    # Re-load and compare keymaps + data — semantic preservation.
+    keymap2 = {}
+    data2 = nt.loads(out, top="any", keymap=keymap2)
+    assert data == data2
+    # All four comment kinds for ('key',) survive:
+    loc = keymap2[("key",)]
+    assert [c.text for c in loc.get_key_leading_comments()]   == ["leading comment for key"]
+    assert [c.text for c in loc.get_key_trailing_comments()]  == ["trailing comment for key"]
+    assert [c.text for c in loc.get_value_leading_comments()] == ["leading comment for value"]
+    assert [c.text for c in loc.get_value_trailing_comments()] == ["trailing comment for value"]
+
+
+def test_force_multiline_for_dict_value_with_kt_or_vl():
+    """If the value is a dict/list and key_trailing/value_leading
+    comments exist, the existing multi-line form is used (regression
+    check that the new branch doesn't break collection-valued items)."""
+    src = (
+        "outer:\n"
+        "        # trailing on outer's key\n"
+        "    inner: 1\n"
+    )
+    keymap = {}
+    data = nt.loads(src, top="any", keymap=keymap)
+    out = nt.dumps(data, map_keys=keymap)
+    keymap2 = {}
+    nt.loads(out, top="any", keymap=keymap2)
+    assert [c.text for c in keymap2[("outer",)].get_key_trailing_comments()] == [
+        "trailing on outer's key"
+    ]
+
+
+def test_force_multiline_value_for_non_string_scalar():
+    """A non-string scalar with a key_trailing comment is still emitted
+    in the multi-line value form."""
+    keymap = {}
+    annotate(("count",), keymap, key_trailing=[Comment("must be int")])
+    data = {"count": 42}
+    out = nt.dumps(data, map_keys=keymap)
+    # The 42 must end up on its own line (with `> ` leader), not inline.
+    assert "count: 42" not in out
+    assert "count:" in out
+    assert "> 42" in out
+    assert "# must be int" in out
+
+
+def test_force_multiline_with_parent_provider():
+    """A parent's key_trailing provider also forces multi-line value
+    form for its children, even when the provider's return value is
+    empty for the specific child."""
+    def kt(k):
+        return [Comment(f"trail {k}")] if k == "verbose" else []
+    keymap = {}
+    annotate((), keymap, key_trailing=kt)
+    data = {"verbose": "yes", "quiet": "no"}
+    out = nt.dumps(data, map_keys=keymap)
+    # verbose: comment lands; the value is on its own line.
+    assert "# trail verbose" in out
+    # 'verbose' and 'quiet' both rendered in multi-line form (provider
+    # presence forces it conservatively).
+    assert "verbose: yes" not in out
+    assert "> yes" in out
+    assert "> no" in out
+
+
 def test_loader_auto_blank_still_fires():
     """Two same-indent loader-built Comments still get an auto-blank
     between them, so the boundary survives a re-load."""