Edit and dump TOML with tomlkit while preserving comments and order: working versions (tested 2026-10-15)

tomlkit==0.15.1
pytest==9.1.1

from __future__ import annotations

import tomlkit
from tomlkit import TOMLDocument


def load(text: str) -> TOMLDocument:
    return tomlkit.parse(text)


def dump(doc: TOMLDocument) -> str:
    return tomlkit.dumps(doc)


def bump_version(doc: TOMLDocument, version: str) -> None:
    doc["project"]["version"] = version


def add_dependency(doc: TOMLDocument, dep: str) -> None:
    doc["project"]["dependencies"].append(dep)


def update_section(doc: TOMLDocument, name: str, values: dict) -> None:
    section = doc.setdefault(name, tomlkit.table())
    for key, value in values.items():
        section[key] = value


def add_note(doc: TOMLDocument, text: str) -> None:
    doc.add(tomlkit.comment(text))
    doc.add(tomlkit.nl())

from config import (
    add_dependency,
    add_note,
    bump_version,
    dump,
    load,
    update_section,
)

SOURCE = """\
# Managed by hand, keep the comments
[project]
name = "demo" # the distribution name
version = "1.0.0"
dependencies = [
    "httpx", # pinned upstream
    "pydantic",
]

[tool.uv]
dev = ["pytest"]
"""


def test_dump_round_trips_comments_and_order():
    assert dump(load(SOURCE)) == SOURCE


def test_bump_version_keeps_inline_comment():
    doc = load(SOURCE)
    bump_version(doc, "1.1.0")
    out = dump(doc)
    assert 'version = "1.1.0"' in out
    assert 'name = "demo" # the distribution name' in out


def test_append_dependency_keeps_array_style():
    doc = load(SOURCE)
    add_dependency(doc, "tomlkit")
    out = dump(doc)
    assert '"httpx", # pinned upstream' in out
    assert '"tomlkit",\n]' in out
    assert doc["project"]["dependencies"].unwrap() == ["httpx", "pydantic", "tomlkit"]


def test_replacing_a_table_naively_drops_comments():
    doc = load(SOURCE)
    doc["project"] = {"name": "demo", "version": "2.0.0"}
    out = dump(doc)
    assert "the distribution name" not in out
    assert "dependencies" not in out


def test_update_section_in_place_keeps_comments():
    doc = load(SOURCE)
    update_section(doc, "project", {"version": "2.0.0", "license": "MIT"})
    out = dump(doc)
    assert 'name = "demo" # the distribution name' in out
    assert 'license = "MIT"' in out
    assert '"httpx", # pinned upstream' in out


def test_add_note_appends_a_standalone_comment():
    doc = load(SOURCE)
    add_note(doc, "generated by ProvenPin")
    assert dump(doc).rstrip().endswith("# generated by ProvenPin")

[pytest]
pythonpath = .