diff --git a/tests/test_api.py b/tests/test_api.py index 14d245b6..5ee42d62 100644 --- a/tests/test_api.py +++ b/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() diff --git a/tomlkit/_utils.py b/tomlkit/_utils.py index 8721f030..d91c8309 100644 --- a/tomlkit/_utils.py +++ b/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) diff --git a/tomlkit/api.py b/tomlkit/api.py index 3582e8ee..d17324e7 100644 --- a/tomlkit/api.py +++ b/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)) diff --git a/tomlkit/items.py b/tomlkit/items.py index 4ca1b1ec..55c52ebc 100644 --- a/tomlkit/items.py +++ b/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: