Fix array item removal (#409)

Author: DimmaDont

tests/test_items.py

@@ -1032,3 +1032,56 @@ 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
+
+    # comment
+'1' #c
+    , #d
+# another comment
+'2'
+ # yet another comment
+,'3', # f
+ # comments here
+'4'#g
+#    comments there
+ ,'5' #h
+    #   comments everywhere
+ # so many comments
+,'6' ,
+      # you get a comment
+# you get a comment
+   '7' ,#j
+# everybody gets a comment!!!
+ "8"#c
+,"9",  \n  "10"
+]
+"""
+    assert expected == parse(expected).as_string()
+    for i in range(11):
+        doc = parse(expected)
+        x = doc["x"]
+        assert isinstance(x, Array)
+        x.remove(str(i))
+        parse(doc.as_string())
+
+
+def test_array_item_removal_newline_restore_next():
+    expected = "x = [\n  '0',\n    '2'\n]"
+
+    docstr = "x = [\n  '0',\n    '1','2'\n]"
+    doc = parse(docstr)
+    doc["x"].remove("1")
+    assert doc.as_string() == expected
+    parse(doc.as_string())
+
+    docstr = "x = [\n  '0',\n    '1',  '2'\n]"
+    doc = parse(docstr)
+    doc["x"].remove("1")
+    assert doc.as_string() == expected
+    parse(doc.as_string())

tomlkit/items.py

@@ -1138,11 +1138,13 @@ def _group_values(self, value: list[Item]) -> list[_ArrayItemGroup]:
         """Group the values into (indent, value, comma, comment) tuples"""
         groups = []
         this_group = _ArrayItemGroup()
+        start_new_group = False
         for item in value:
             if isinstance(item, Whitespace):
-                if "," not in item.s:
+                if "," not in item.s or start_new_group:
                     groups.append(this_group)
                     this_group = _ArrayItemGroup(indent=item)
+                    start_new_group = False
                 else:
                     if this_group.value is None:
                         # when comma is met and no value is provided, add a dummy Null
@@ -1152,6 +1154,8 @@ def _group_values(self, value: list[Item]) -> list[_ArrayItemGroup]:
                 if this_group.value is None:
                     this_group.value = Null()
                 this_group.comment = item
+                # Comments are the last item in a group.
+                start_new_group = True
             elif this_group.value is None:
                 this_group.value = item
             else:
@@ -1399,6 +1403,7 @@ def __delitem__(self, key: int | slice):
                 if not isinstance(key, slice):
                     raise IndexError("list index out of range") from e
             else:
+                group_rm = self._value[idx]
                 del self._value[idx]
                 if (
                     idx == 0
@@ -1408,6 +1413,44 @@ def __delitem__(self, key: int | slice):
                 ):
                     # Remove the indentation of the first item if not newline
                     self._value[idx].indent = None
+                comma_in_indent = (
+                    group_rm.indent is not None and "," in group_rm.indent.s
+                )
+                comma_in_comma = group_rm.comma is not None and "," in group_rm.comma.s
+                if comma_in_indent and comma_in_comma:
+                    # Removed group had both commas. Add one to the next group.
+                    group = self._value[idx] if len(self._value) > idx else None
+                    if group is not None:
+                        if group.indent is None:
+                            group.indent = Whitespace(",")
+                        elif "," not in group.indent.s:
+                            # Insert the comma after the newline
+                            try:
+                                newline_index = group.indent.s.index("\n")
+                                group.indent._s = (
+                                    group.indent.s[: newline_index + 1]
+                                    + ","
+                                    + group.indent.s[newline_index + 1 :]
+                                )
+                            except ValueError:
+                                group.indent._s = "," + group.indent.s
+                elif not comma_in_indent and not comma_in_comma:
+                    # Removed group had no commas. Remove the next comma found.
+                    for j in range(idx, len(self._value)):
+                        group = self._value[j]
+                        if group.indent is not None and "," in group.indent.s:
+                            group.indent._s = group.indent.s.replace(",", "", 1)
+                            break
+                if group_rm.indent is not None and "\n" in group_rm.indent.s:
+                    # Restore the removed group's newline onto the next group
+                    # if the next group does not have a newline.
+                    # i.e. the two were on the same line
+                    group = self._value[idx] if len(self._value) > idx else None
+                    if group is not None and (
+                        group.indent is None or "\n" not in group.indent.s
+                    ):
+                        group.indent = group_rm.indent
+
         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
+                if isinstance(elems[-1], Whitespace):
+                    elems[-1]._s = elems[-1].s + ","
                 else:
-                    comma = Whitespace(",")
-                elems.append(comma)
+                    elems.append(Whitespace(","))
                 prev_value = False
                 continue