Preserve literal quotes when replacing table strings

CHANGELOG.md

@@ -21,6 +21,7 @@
 
 ### Fixed
 
+- Preserve single-line literal quotes when replacing a table field with a plain string, falling back to basic quoting when the value requires escaping. ([#325](https://github.com/python-poetry/tomlkit/issues/325))
 - Fix `string()` dropping a leading newline of a multiline string on round-trip: a value beginning with a newline is now rendered with an extra leading newline (the one the parser trims after the opening delimiter) so it survives re-parsing.
 - Fix invalid serialization with a duplicated comma when removing a non-edge element from a parsed inline table. ([#486](https://github.com/python-poetry/tomlkit/pull/486))
 - Fix invalid serialization with a duplicated comma when appending or inserting into a comma-first formatted array. ([#499](https://github.com/python-poetry/tomlkit/pull/499))

tests/test_toml_document.py

@@ -22,6 +22,74 @@
 from tomlkit.toml_document import TOMLDocument
 
 
+@pytest.mark.parametrize(
+    "content,path",
+    [
+        ("before = 1\n  value  =  'old' # keep\nafter = 2\n", ()),
+        ("before = 1\r\nvalue = 'old' # keep\r\nafter = 2\r\n", ()),
+        ("[table]\nbefore = 1\n  value  =  'old' # keep\nafter = 2\n", ("table",)),
+        ("table = {before = 1, value = 'old', after = 2}\n", ("table",)),
+        (
+            "[table.a]\nx = 1\n[other]\ny = 2\n[table]\nvalue = 'old' # keep\n",
+            ("table",),
+        ),
+    ],
+)
+def test_replace_literal_string_preserves_table_style(
+    content: str, path: tuple[str, ...]
+) -> None:
+    doc = parse(content)
+    table: Any = doc
+    for key in path:
+        table = table[key]
+    table["value"] = "new"
+    assert doc.as_string() == content.replace("'old'", "'new'")
+    assert table["value"] == dict(table)["value"] == table.unwrap()["value"] == "new"
+    assert parse(doc.as_string()).unwrap() == doc.unwrap()
+
+    # Explicit items retain the caller's quote choice instead of inheriting it.
+    table["value"] = tomlkit.string("explicit")
+    assert doc.as_string() == content.replace("'old'", '"explicit"')
+    table["value"] = tomlkit.string("literal", literal=True)
+    assert doc.as_string() == content.replace("'old'", "'literal'")
+    table.update(value="updated")
+    assert doc.as_string() == content.replace("'old'", "'updated'")
+    table["value"] = "can't"
+    assert doc.as_string() == content.replace("'old'", '"can\'t"')
+    assert parse(doc.as_string()).unwrap() == doc.unwrap()
+
+
+@pytest.mark.parametrize("value", ["", "unicode λ", "a\\b", 'a"b', "a\tb"])
+def test_replace_literal_string_preserves_value(value: str) -> None:
+    doc = parse("value = 'old'\n")
+    doc["value"] = value
+    assert doc.as_string() == f"value = '{value}'\n"
+    assert parse(doc.as_string())["value"] == value
+
+
+@pytest.mark.parametrize("value", ["can't", "a\nb", "a\rb", "a\x00b", "a\x7fb"])
+def test_replace_literal_string_falls_back_to_basic(value: str) -> None:
+    doc = parse("value = 'old' # keep\n")
+    doc["value"] = value
+    assert doc.as_string() == f"value = {tomlkit.string(value).as_string()} # keep\n"
+    assert parse(doc.as_string())["value"] == value
+
+
+@pytest.mark.parametrize("original", ['"old"', "'''old'''", '"""old"""', "1"])
+def test_replace_other_items_keeps_default_string_style(original: str) -> None:
+    doc = parse(f"value = {original}\n")
+    doc["value"] = "new"
+    assert doc.as_string() == 'value = "new"\n'
+
+
+def test_literal_replacement_does_not_change_new_keys_or_arrays() -> None:
+    doc = parse("value = 'old'\narray = ['old']\n")
+    doc["new"] = "new"
+    doc["array"][0] = "new"
+    doc["value"] = 2
+    assert doc.as_string() == 'value = 2\narray = ["new"]\nnew = "new"\n'
+
+
 def test_document_is_a_dict(example: Callable[[str], str]) -> None:
     content = example("example")
 

tomlkit/container.py

@@ -14,6 +14,7 @@
 from tomlkit._compat import decode
 from tomlkit._types import _CustomDict
 from tomlkit._utils import merge_dicts
+from tomlkit.exceptions import InvalidStringError
 from tomlkit.exceptions import KeyAlreadyPresent
 from tomlkit.exceptions import NonExistentKey
 from tomlkit.exceptions import TOMLKitError
@@ -23,6 +24,8 @@
 from tomlkit.items import Key
 from tomlkit.items import Null
 from tomlkit.items import SingleKey
+from tomlkit.items import String
+from tomlkit.items import StringType
 from tomlkit.items import Table
 from tomlkit.items import Trivia
 from tomlkit.items import Whitespace
@@ -847,9 +850,26 @@ def __contains__(self, key: object) -> bool:
                 OutOfOrderTableProxy(self, idx)
         return True
 
+    def _item_for_key(
+        self, key: Key | str, value: object, parent: Item | None = None
+    ) -> Item:
+        # Inspect plain strings before conversion, so an explicit String keeps
+        # the caller's chosen quoting style.
+        if isinstance(value, str) and not isinstance(value, String):
+            previous = self.get(key)
+            if isinstance(previous, String) and previous._t is StringType.SLL:
+                try:
+                    return String.from_raw(value, StringType.SLL)
+                except InvalidStringError:
+                    # Some values need the escaping provided by a basic string.
+                    pass
+        return _item(value, _parent=parent)
+
     def __setitem__(self, key: Key | str, value: Any) -> None:
         if key in self:
             old_key = next(filter(lambda k: k == key, self._map))
+            if isinstance(value, str) and not isinstance(value, String):
+                value = self._item_for_key(key, value)
             self._replace(old_key, key, value)
         else:
             self.append(key, value)

tomlkit/items.py

@@ -1854,7 +1854,7 @@ def __contains__(self, key: object) -> bool:
 
     def __setitem__(self, key: Key | str, value: Any) -> None:
         if not isinstance(value, Item):
-            value = item(value, _parent=self)
+            value = self._value._item_for_key(key, value, parent=self)
 
         is_replace = key in self
         self._value[key] = value
@@ -2198,7 +2198,7 @@ def __setitem__(self, key: Key | str, value: Any) -> None:
         if hasattr(value, "trivia") and value.trivia.comment:
             value.trivia.comment = ""
         if not isinstance(value, Item):
-            value = item(value, _parent=self)
+            value = self._value._item_for_key(key, value, parent=self)
         self._validate_child(value)
         super().__setitem__(key, value)