Keep comments around arrays of tables

CHANGELOG.md

@@ -4,6 +4,13 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+### Fixed
+
+- A comment before the third or later element of an array of tables is no longer moved after the first element and then lost on the next run
+- Comments before a nested array of tables (`[[a.b]]` inside `[a]`) are no longer dropped, and the parent table's comment stays on the parent
+
 ## 0.24.4
 
 ### Changed

tests/examples/aot-comments.toml

@@ -0,0 +1,13 @@
+# A comment before an element of an array of tables stays with that
+# element, including the third and later ones.
+
+[[fruit]]
+name = "apple"
+
+# The second fruit
+[[fruit]]
+name = "banana"
+
+# The third fruit
+[[fruit]]
+name = "cherry"

tests/examples/aot-nested-comments.toml

@@ -0,0 +1,29 @@
+# Comments on a table, on a nested array of tables, and on arrays of
+# tables under implicit super tables all survive.
+
+[a]
+y = 1
+
+# About m
+[m]
+# About the first n
+[[m.n]]
+x = 1
+
+# About the second n
+[[m.n]]
+x = 2
+
+# Reached through the implicit super tables p and p.q
+[[p.q.r]]
+z = 1
+
+[[s]]
+# First t
+[[s.t]]
+x = 1
+
+[[s]]
+# Second t
+[[s.t]]
+x = 2

tests/examples/sorted/aot-comments.toml

@@ -0,0 +1,13 @@
+# A comment before an element of an array of tables stays with that
+# element, including the third and later ones.
+
+[[fruit]]
+name = "apple"
+
+# The second fruit
+[[fruit]]
+name = "banana"
+
+# The third fruit
+[[fruit]]
+name = "cherry"

tests/examples/sorted/aot-nested-comments.toml

@@ -0,0 +1,32 @@
+# Comments on a table, on a nested array of tables, and on arrays of
+# tables under implicit super tables all survive.
+
+[a]
+y = 1
+
+# About m
+[m]
+
+# About the first n
+[[m.n]]
+x = 1
+
+# About the second n
+[[m.n]]
+x = 2
+
+# Reached through the implicit super tables p and p.q
+[[p.q.r]]
+z = 1
+
+[[s]]
+
+# First t
+[[s.t]]
+x = 1
+
+[[s]]
+
+# Second t
+[[s.t]]
+x = 2

tests/test_toml_sort.py

@@ -43,6 +43,16 @@ def test_sort_toml_is_str() -> None:
                 ),
             },
         ),
+        (
+            "aot-comments",
+            "aot-comments",
+            {},
+        ),
+        (
+            "aot-nested-comments",
+            "aot-nested-comments",
+            {},
+        ),
         (
             "inline",
             "inline",

toml_sort/tomlsort.py

@@ -589,11 +589,11 @@ def toml_elements_sorted(
                 self.format_config.spaces_before_inline_comment,
             )
             for table in original.children:
-                previous_item = next(iter(new_aot), parent)
+                # Comments belong after the element that precedes them,
+                # which is the last one appended so far.
+                previous_item = new_aot[-1] if len(new_aot) else parent
                 attach_comments(table, previous_item)
-                new_aot.append(
-                    self.toml_elements_sorted(table, next(iter(new_aot), previous_item))
-                )
+                new_aot.append(self.toml_elements_sorted(table, previous_item))
 
             return new_aot
 
@@ -728,23 +728,17 @@ def table_to_tomlsortitem(
         if not value.is_super_table():
             new_table.trivia.indent = "\n"
 
-        first_child = next(iter(children), None)
-
-        # If the first child of this item is an AoT, we want the
-        # comment to be attached to the first table within the AoT,
-        # rather than the parent AoT object
-        if first_child and first_child.is_aot:
-            first_child.children[0].attached_comments = comments
-            comments = []
-
-        # If this item is a super table we want to walk down
-        # the tree and attach the comment to the first non-super table.
-        if value.is_super_table():
+        # A super table has no header of its own, so its comment must be
+        # attached to the first thing that is rendered beneath it: walk down
+        # through nested super tables and arrays of tables to the first real
+        # table. Prepend rather than replace, so comments that were already
+        # attached to that table are kept.
+        if value.is_super_table() and children:
             child_table = children[0]
-            while child_table.is_super_table:
+            while child_table.is_super_table or child_table.is_aot:
                 child_table = child_table.children[0]
 
-            child_table.attached_comments = comments
+            child_table.attached_comments = comments + child_table.attached_comments
             comments = []
 
         item = TomlSortItem(keys, new_table, comments, children)