Fix array item grouping and removal

Author: DimmaDont

tests/test_items.py

@@ -1019,3 +1019,29 @@ def test_removal_of_arrayitem_with_extra_whitespace():
     docstr = doc.as_string()
     parse(docstr)
     assert docstr == expected
+
+
+def test_badly_formatted_array_and_item_removal():
+    expected = """
+x = [
+    '0'#a
+,#b
+'1' #c
+    , #d
+'2' # e
+,'3', # f
+'4'#g
+ ,'5' #h
+,'6' ,
+   '7' ,#j
+'8' , #k
+   '9' , # l
+]
+"""
+    assert expected == parse(expected).as_string()
+    for i in range(10):
+        doc = parse(expected)
+        x = doc["x"]
+        assert isinstance(x, Array)
+        x.remove(str(i))
+        parse(doc.as_string())

tomlkit/items.py

@@ -1083,30 +1083,33 @@ def _getstate(self, protocol: int = 3) -> tuple:
 
 
 class _ArrayItemGroup:
-    __slots__ = ("comma", "comment", "indent", "value")
+    __slots__ = ("comma", "comment", "comment2", "indent", "value")
 
     def __init__(
         self,
         value: Item | None = None,
         indent: Whitespace | None = None,
         comma: Whitespace | None = None,
         comment: Comment | None = None,
+        comment2: Comment | None = None,
     ) -> None:
         self.value = value
         self.indent = indent
         self.comma = comma
         self.comment = comment
+        self.comment2 = comment2
 
     def __iter__(self) -> Iterator[Item]:
         return filter(
-            lambda x: x is not None, (self.indent, self.value, self.comma, self.comment)
+            lambda x: x is not None,
+            (self.indent, self.value, self.comment, self.comma, self.comment2),
         )
 
     def __repr__(self) -> str:
         return repr(tuple(self))
 
     def is_whitespace(self) -> bool:
-        return self.value is None and self.comment is None
+        return self.value is None and self.comment is None and self.comment2 is None
 
     def __bool__(self) -> bool:
         try:
@@ -1135,10 +1138,10 @@ def __init__(
         self._reindex()
 
     def _group_values(self, value: list[Item]) -> list[_ArrayItemGroup]:
-        """Group the values into (indent, value, comma, comment) tuples"""
+        """Group the values into (indent, value, comment, comma, comment2) tuples"""
         groups = []
         this_group = _ArrayItemGroup()
-        for item in value:
+        for i, item in enumerate(value):
             if isinstance(item, Whitespace):
                 if "," not in item.s:
                     groups.append(this_group)
@@ -1151,7 +1154,16 @@ def _group_values(self, value: list[Item]) -> list[_ArrayItemGroup]:
             elif isinstance(item, Comment):
                 if this_group.value is None:
                     this_group.value = Null()
-                this_group.comment = item
+                # It's possible to have two comments per group.
+                # Comment one, if present, is always immediately after the value,
+                # and Comment two after a Whitespace with comma.
+                prev_item = value[i - 1]
+                if isinstance(prev_item, Whitespace) and "," in prev_item.s:
+                    # If this is Comment two, this is also the end of the current group,
+                    # but let that be done in the next pass
+                    this_group.comment2 = item
+                else:
+                    this_group.comment = item
             elif this_group.value is None:
                 this_group.value = item
             else:
@@ -1209,8 +1221,9 @@ def as_string(self) -> str:
             self.trivia.indent
             + " " * 4
             + v.value.as_string()
-            + ("," if not isinstance(v.value, Null) else "")
             + (v.comment.as_string() if v.comment is not None else "")
+            + ("," if not isinstance(v.value, Null) else "")
+            + (v.comment2.as_string() if v.comment2 is not None else "")
             + "\n"
             for v in self._value
             if v.value is not None
@@ -1372,6 +1385,11 @@ def insert(self, pos: int, value: Any) -> None:
                 # Copy the comma from the last item if 1) it contains a value and
                 # 2) the array is multiline
                 comma = last_item.comma
+            if last_item.comment and not last_item.comma:
+                # Arrays with a single value might not have a comma after the value, but may have a comment.
+                # Move comment one to comment two if so.
+                last_item.comment2 = last_item.comment
+                last_item.comment = None
             if last_item.comma is None and not isinstance(last_item.value, Null):
                 # Add comma to the last item to separate it from the following items.
                 last_item.comma = Whitespace(",")
@@ -1408,6 +1426,14 @@ def __delitem__(self, key: int | slice):
                 ):
                     # Remove the indentation of the first item if not newline
                     self._value[idx].indent = None
+                next_group = self._value[idx] if len(self._value) > idx else None
+                # If the previous group ends with a comment,
+                if self._value[idx - 1].comment2 and next_group:
+                    # and the next group doesn't start with a newline, add/set one
+                    if next_group.indent and "\n" not in next_group.indent.s:
+                        next_group.indent = Whitespace("\n" + next_group.indent.s)
+                    else:
+                        next_group.indent = Whitespace("\n")
         if len(self._value) > 0:
             v = self._value[-1]
             if not v.is_whitespace():

tomlkit/parser.py

@@ -595,15 +595,11 @@ def _parse_array(self) -> Array:
             # consume comma
             if prev_value and self._current == ",":
                 self.inc(exception=UnexpectedEofError)
-                # Check if the previous item is Whitespace
-                if isinstance(elems[-1], Whitespace) and " " in elems[-1].s:
-                    # Preserve the previous whitespace
-                    comma = Whitespace(elems[-1].s + ",")
-                    # Remove the replaced item
-                    del elems[-1]
+                # If the previous item is Whitespace, add to it instead of appending a new Whitespace
+                if isinstance(elems[-1], Whitespace):
+                    elems[-1] = Whitespace(elems[-1].s + ",")
                 else:
-                    comma = Whitespace(",")
-                elems.append(comma)
+                    elems.append(Whitespace(","))
                 prev_value = False
                 continue