Reject forbidden control characters in generated comments

tests/test_api.py

@@ -292,6 +292,54 @@ def test_comment_multiline_is_valid_toml() -> None:
     assert parse(rendered).as_string() == rendered
 
 
+@pytest.mark.parametrize(
+    "character", [chr(code) for code in range(32) if code not in (9, 10, 13)] + ["\x7f"]
+)
+@pytest.mark.parametrize("kind", ["standalone", "item", "array"])
+def test_comments_reject_control_characters(character: str, kind: str) -> None:
+    doc = parse("value = 1 # original\nvalues = [1]\n")
+    original = doc.as_string()
+    text = f"before{character}after"
+
+    with pytest.raises(ValueError, match="Comment cannot contain"):
+        if kind == "standalone":
+            doc.add(tomlkit.comment(text))
+        elif kind == "item":
+            doc["value"].comment(text)
+        else:
+            doc["values"].add_line(2, comment=text)
+
+    assert doc.as_string() == original
+    assert doc.unwrap() == {"value": 1, "values": [1]}
+
+
+@pytest.mark.parametrize("text", ["before\rafter", "before\r", "before\r\r\nafter"])
+def test_comment_rejects_bare_carriage_return(text: str) -> None:
+    with pytest.raises(ValueError, match="Comment cannot contain"):
+        tomlkit.comment(text)
+
+
+@pytest.mark.parametrize("newline", ["\n", "\r\n"])
+def test_comment_preserves_valid_line_endings(newline: str) -> None:
+    doc = tomlkit.document()
+    doc.add(tomlkit.comment(f"one{newline}two{newline}"))
+    rendered = f"# one{newline}# two{newline}#\n"
+    assert doc.as_string() == rendered
+    assert parse(rendered).as_string() == rendered
+
+
+def test_comment_apis_preserve_tabs_and_unicode() -> None:
+    text = "tab\tunicode \u00e9 \u0085 \u2028"
+    doc = parse("value = 1\nvalues = []\n")
+    doc.add(tomlkit.comment(text))
+    doc["value"].comment(text)
+    doc["values"].add_line(2, comment=text)
+    rendered = doc.as_string()
+    assert rendered.count(text) == 3
+    assert parse(rendered).as_string() == rendered
+    assert parse(rendered).unwrap() == {"value": 1, "values": [2]}
+
+
 def test_array() -> None:
     a = tomlkit.array()
 

tomlkit/_utils.py

@@ -123,6 +123,15 @@ def parse_rfc3339(string: str) -> datetime | date | time:
 _basic_escapes = CONTROL_CHARS | {'"', "\\"}
 
 
+def validate_comment(comment: str, *, multiline: bool = False) -> None:
+    """Reject forbidden controls, optionally allowing LF and CRLF line endings."""
+    if not multiline and ("\n" in comment or "\r" in comment):
+        raise ValueError("Comment cannot contain line breaks")
+    for char in comment.replace("\r\n", "\n"):
+        if char in CONTROL_CHARS and char not in "\t\n":
+            raise ValueError(f"Comment cannot contain control character {char!r}")
+
+
 def _unicode_escape(seq: str) -> str:
     return "".join(f"\\u{ord(c):04x}" for c in seq)
 

tomlkit/api.py

@@ -12,6 +12,7 @@
 from typing import TypeVar
 
 from tomlkit._utils import parse_rfc3339
+from tomlkit._utils import validate_comment
 from tomlkit.container import Container
 from tomlkit.exceptions import UnexpectedCharError
 from tomlkit.items import CUSTOM_ENCODERS
@@ -349,6 +350,7 @@ def comment(string: str) -> Comment:
     A multiline string produces one ``#``-prefixed line per line so that the
     result is still valid TOML.
     """
+    validate_comment(string, multiline=True)
     lines = string.split("\n")
     rendered = "\n".join(f"# {line}" if line else "#" for line in lines)
     return Comment(Trivia(comment_ws="  ", comment=rendered))

tomlkit/items.py

@@ -31,6 +31,7 @@
 from tomlkit._types import _CustomList
 from tomlkit._utils import CONTROL_CHARS
 from tomlkit._utils import escape_string
+from tomlkit._utils import validate_comment
 from tomlkit.exceptions import ConvertError
 from tomlkit.exceptions import InvalidStringError
 
@@ -575,8 +576,7 @@ def unwrap(self) -> Any:
 
     def comment(self, comment: str) -> Item:
         """Attach a comment to this item"""
-        if "\n" in comment or "\r" in comment:
-            raise ValueError("Comment cannot contain line breaks")
+        validate_comment(comment)
         if not comment.strip().startswith("#"):
             comment = "# " + comment
 
@@ -1618,8 +1618,8 @@ def add_line(
             4, 5, 6,
         ]
         """
-        if comment and ("\n" in comment or "\r" in comment):
-            raise ValueError("Comment cannot contain line breaks")
+        if comment:
+            validate_comment(comment)
         new_values: list[Item] = []
         first_indent = f"\n{indent}" if newline else indent
         if first_indent: