Preserve whitespace when replacing super tables

tests/test_toml_document.py

@@ -965,6 +965,24 @@ def test_replace_preserve_sep() -> None:
     )
 
 
+def test_replace_super_table_preserves_whitespace() -> None:
+    content = """\
+[env.pro1.rst]
+name = "x7"
+
+[env2]
+name = 2
+
+[env3]
+name = 3
+"""
+    doc = parse(content)
+
+    doc["env"] = doc["env"]
+
+    assert doc.as_string() == content
+
+
 def test_replace_with_table_of_nested() -> None:
     example = """\
     [a]

tomlkit/container.py

@@ -1152,9 +1152,12 @@ def ends_with_whitespace(it: Any) -> bool:
     """Returns ``True`` if the given item ``it`` is a ``Table`` or ``AoT`` object
     ending with a ``Whitespace``.
     """
-    return (
-        isinstance(it, Table) and isinstance(it.value._previous_item(), Whitespace)
-    ) or (isinstance(it, AoT) and len(it) > 0 and isinstance(it[-1], Whitespace))
+    if isinstance(it, Whitespace):
+        return True
+    if isinstance(it, Table):
+        previous = it.value._previous_item()
+        return previous is not None and ends_with_whitespace(previous)
+    return isinstance(it, AoT) and len(it) > 0 and ends_with_whitespace(it[-1])
 
 
 def _equal_with_nan(left: Any, right: Any) -> bool: