Repository: sdispater/tomlkit

Issue:
style not preserved when modifying dict
First off, thanks for this library. A style-preserving writer is a really useful and rare thing. Here's my config:

```toml
[versions]
robotpy-wpiutil             = "2022.2.1.1"
pyntcore                    = "2022.2.1.1"
robotpy-hal                 = "2022.2.1.1"
robotpy-wpimath             = "2022.2.1.2"
wpilib                      = "2022.2.1.3"
robotpy-halsim-gui          = "2022.2.1.0"
robotpy-halsim-ds-socket    = "2022.2.1.0"
robotpy-halsim-ws           = "2022.2.1.0"
robotpy-commands-v1         = "2022.2.1.1"
robotpy-commands-v2         = "2022.2.1.1"


```

When I run

```python
import tomlkit

with open("cfg.toml") as fp:
    cfgdata = tomlkit.parse(fp.read())

cfgdata["versions"]["robotpy-wpimath"] = "2022.2.1.4"
print(tomlkit.dumps(cfgdata))
```

The output is

```toml
[versions]
robotpy-wpiutil             = "2022.2.1.1"
pyntcore                    = "2022.2.1.1"
robotpy-hal                 = "2022.2.1.1"
robotpy-wpimath = "2022.2.1.4"
wpilib                      = "2022.2.1.3"
robotpy-halsim-gui          = "2022.2.1.0"
robotpy-halsim-ds-socket    = "2022.2.1.0"
robotpy-halsim-ws           = "2022.2.1.0"
robotpy-commands-v1         = "2022.2.1.1"
robotpy-commands-v2         = "2022.2.1.1"
```
For now, my workaround is to directly assign the _value, and the style information is retained:

```python
import tomlkit

with open("cfg.toml") as fp:
    cfgdata = tomlkit.parse(fp.read())

cfgdata["versions"]["robotpy-wpimath"]._value = "2022.2.1.4"
print(tomlkit.dumps(cfgdata))
```

Is this intended behavior, or am I missing something? From reading the quickstart, it seemed to me that just assigning the dictionary would keep style information, but that doesn't seem to be the case.


Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5ed9ec7..24fc8f8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## [Unreleased]
 
+### Fixed
+
+- Fix a bug of separators not being kept when replacing the value. ([#170](https://github.com/sdispater/tomlkit/pull/170))
+
 ## [0.9.0] - 2022-02-01
 
 ### Added
diff --git a/tomlkit/container.py b/tomlkit/container.py
index c6a4d8a..c20a6ab 100644
--- a/tomlkit/container.py
+++ b/tomlkit/container.py
@@ -619,7 +619,8 @@ class Container(_CustomDict):
 
     def __setitem__(self, key: Union[Key, str], value: Any) -> None:
         if key is not None and key in self:
-            self._replace(key, key, value)
+            old_key = next(filter(lambda k: k == key, self._map))
+            self._replace(old_key, key, value)
         else:
             self.append(key, value)
 
@@ -636,9 +637,6 @@ class Container(_CustomDict):
         if not isinstance(key, Key):
             key = SingleKey(key)
 
-        if not isinstance(new_key, Key):
-            new_key = SingleKey(new_key)
-
         idx = self._map.get(key, None)
         if idx is None:
             raise NonExistentKey(key)
@@ -648,8 +646,7 @@ class Container(_CustomDict):
     def _replace_at(
         self, idx: Union[int, Tuple[int]], new_key: Union[Key, str], value: Item
     ) -> None:
-        if not isinstance(new_key, Key):
-            new_key = SingleKey(new_key)
+        value = _item(value)
 
         if isinstance(idx, tuple):
             for i in idx[1:]:
@@ -658,17 +655,21 @@ class Container(_CustomDict):
             idx = idx[0]
 
         k, v = self._body[idx]
+        if not isinstance(new_key, Key):
+            if (
+                isinstance(value, (AoT, Table)) != isinstance(v, (AoT, Table))
+                or new_key != k.key
+            ):
+                new_key = SingleKey(new_key)
+            else:  # Inherit the sep of the old key
+                new_key = k
 
-        self._map[new_key] = self._map.pop(k)
+        del self._map[k]
+        self._map[new_key] = idx
         if new_key != k:
             dict.__delitem__(self, k)
 
-        if isinstance(self._map[new_key], tuple):
-            self._map[new_key] = self._map[new_key][0]
-
-        value = _item(value)
-
-        if isinstance(value, (AoT, Table)) and not isinstance(v, (AoT, Table)):
+        if isinstance(value, (AoT, Table)) != isinstance(v, (AoT, Table)):
             # new tables should appear after all non-table values
             self.remove(k)
             for i in range(idx, len(self._body)):


Test patch:
diff --git a/tests/test_toml_document.py b/tests/test_toml_document.py
index 5445134..8d79ed4 100644
--- a/tests/test_toml_document.py
+++ b/tests/test_toml_document.py
@@ -735,6 +735,44 @@ foo = "bar"
     )
 
 
+def test_replace_table_with_value():
+    content = """[foo]
+a = 1
+
+[bar]
+b = 2
+"""
+    doc = parse(content)
+    doc["bar"] = 42
+    assert (
+        doc.as_string()
+        == """bar = 42
+[foo]
+a = 1
+
+"""
+    )
+
+
+def test_replace_preserve_sep():
+    content = """a   =   1
+
+[foo]
+b  =  "what"
+"""
+    doc = parse(content)
+    doc["a"] = 2
+    doc["foo"]["b"] = "how"
+    assert (
+        doc.as_string()
+        == """a   =   2
+
+[foo]
+b  =  "how"
+"""
+    )
+
+
 def test_replace_with_table_of_nested():
     example = """\
     [a]