Add code to store comments during formatting

src/pyproject_fmt/formatter/util.py

@@ -32,24 +32,95 @@ def __gt__(self, __other: Any) -> bool:  # noqa: U101
         ...
 
 
+@dataclass
+class TableComments:
+    key_comments: dict[str, list[Comment]]
+    top_comments: list[Comment]
+    bottom_comments: list[Comment]
+
+
+def store_comments(table: AbstractTable) -> TableComments:
+    """
+    Takes in a table, and returns a `TableComments` object.
+    All comment lines in the table are attached to the key below it.
+    Bottom comments are stored separately.
+
+    The only exception to this rule is: If there's comments at the top of
+    the table, separated from the rest of the body with newlines, then
+    they are stored separately as top comments.
+    """
+    key_comments: dict[str, list[Comment]] = {}
+    top_comments: list[Comment] = []
+    bottom_comments: list[Comment] = []
+
+    body = table.value.body
+    first_key_index = -1
+    for index, (key, _) in enumerate(body):
+        if isinstance(key, Key):
+            first_key_index = index
+            break
+
+    if first_key_index > 0 and isinstance(body[first_key_index - 1], Whitespace):
+        # We may have comments at the top of the file.
+        for _, item in body[:first_key_index]:
+            if isinstance(item, Comment):
+                top_comments.append(item)
+
+        # delete the top comments
+        del body[:first_key_index]
+
+    # Start picking comments and attaching them to each key
+    current_comments: list[Comment] = []
+    for key, item in body:
+        if isinstance(item, Comment):
+            current_comments.append(item)
+        elif isinstance(key, Key):
+            key_comments[key.key] = current_comments
+            current_comments = []
+
+    # If we're left with comments at the bottom
+    if current_comments:
+        bottom_comments = current_comments
+
+    return TableComments(key_comments, top_comments, bottom_comments)
+
+
+def insert_key_with_comments(table: AbstractTable, key: Key, item: Item, comments: TableComments) -> None:
+    body = table.value.body
+    entry_comments = comments.key_comments.get(key.key, [])
+    body.extend((None, comment) for comment in entry_comments)
+    body.append((key, item))
+
+
 def order_keys(
     table: AbstractTable,
     to_pin: Sequence[str] | None = None,
     sort_key: None | Callable[[tuple[str, tuple[Key, Item]]], SupportsDunderLT | SupportsDunderGT] = None,
 ) -> None:
+    comments = store_comments(table)
+
     body = table.value.body
     entries = {i.key: (i, v) for (i, v) in body if isinstance(i, Key)}
     body.clear()
 
+    for comment in comments.top_comments:
+        body.append((None, comment))
+
     for pin in to_pin or []:  # push pinned to start
         if pin in entries:
-            body.append(entries[pin])
+            key, value = entries[pin]
+            insert_key_with_comments(table, key, value, comments)
             del entries[pin]
     # append the rest
     if sort_key is None:
-        body.extend(entries.values())
+        for _, (key, item) in entries.items():
+            insert_key_with_comments(table, key, item, comments)
     else:
-        body.extend(v for k, v in sorted(entries.items(), key=sort_key))
+        for _, (key, item) in sorted(entries.items(), key=sort_key):
+            insert_key_with_comments(table, key, item, comments)
+
+    for comment in comments.bottom_comments:
+        body.append((None, comment))
 
     if isinstance(table, Table):
         body.append((None, Whitespace("\n")))  # add trailing newline to separate

tests/formatter/test_project.py

@@ -105,3 +105,13 @@ def test_entry_points(fmt: Fmt) -> None:
     beta = {C = "c",D = "d"}
     """
     fmt(fmt_project, start, expected)
+
+
+def test_comments(fmt: Fmt) -> None:
+    start = expected = """
+    [project]
+    # A comment-only line
+    name = "something"  # This is the name
+    # Another comment
+    """
+    fmt(fmt_project, start, expected)