Add second comment to _ArrayItemGroup

Author: DimmaDont

tests/test_items.py

@@ -1019,3 +1019,27 @@ def test_removal_of_arrayitem_with_extra_whitespace():
     docstr = doc.as_string()
     parse(docstr)
     assert docstr == expected
+
+
+def test_removal_of_items_from_badly_formatted_arrays():
+    for i in range(10):
+        doc = parse(
+            """
+x = [
+    '0'#a
+,#b
+'1' #c
+    , #d
+'2' # e
+,'3', # f
+'4'#g
+ ,'5' #h
+,'6',#i
+   '7' ,#j
+'8' , #k
+   '9' , # l
+]
+"""
+        )
+        doc["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,12 +1138,25 @@ 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:
+                    # Whitespace without a comma doesn't necessarily mean end of group
+                    # Check the next item to see if it has a comma, which does mean end of group.
+                    if (
+                        len(value) > i + 1
+                        and isinstance(value[i + 1], Whitespace)
+                        and "," in value[i + 1].s
+                    ):
+                        # Prepend this whitespace to that item
+                        value[i + 1] = Whitespace(item.s + value[i + 1].s)
+                        # Set an indent if this group does not have one.
+                        if not this_group.indent:
+                            this_group.indent = Whitespace("\n")
+                        continue
                     groups.append(this_group)
                     this_group = _ArrayItemGroup(indent=item)
                 else:
@@ -1151,7 +1167,13 @@ 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.
+                if isinstance(value[i - 1], Whitespace) and "," in value[i - 1].s:
+                    # If the previous item is Whitespace and contains a comma,
+                    # this must be comment two.
+                    this_group.comment2 = item
+                else:
+                    this_group.comment = item
             elif this_group.value is None:
                 this_group.value = item
             else:
@@ -1209,8 +1231,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 +1395,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(",")