feat: migrate from tomlkit to tomlrt

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: dimbleby

poetry.lock

@@ -21,7 +21,7 @@ dependencies = [
     "requests-toolbelt (>=1.0.0,<2.0.0)",
     "shellingham (>=1.5,<2.0)",
     "tomli (>=2.0.1,<3.0.0) ; python_version < '3.11'",
-    "tomlkit (>=0.11.4,<1.0.0)",
+    "tomlrt (>=1.4,<2)",
     # trove-classifiers uses calver, so version is unclamped
     "trove-classifiers (>=2022.5.19)",
     "virtualenv (>=20.26.6)",

pyproject.toml

@@ -4,8 +4,7 @@
 from typing import TYPE_CHECKING
 from typing import Any
 
-from tomlkit import document
-from tomlkit import table
+from tomlrt import Document
 
 from poetry.config.config_source import ConfigSource
 from poetry.config.config_source import PropertyNotFoundError
@@ -16,8 +15,6 @@
     from collections.abc import Iterator
     from collections.abc import Sequence
 
-    from tomlkit.toml_document import TOMLDocument
-
     from poetry.toml.file import TOMLFile
 
 
@@ -34,37 +31,19 @@ def file(self) -> TOMLFile:
         return self._file
 
     def get_property(self, key: str | Sequence[str]) -> Any:
-        keys = split_key(key)
-
-        config = self.file.read() if self.file.exists() else {}
-
-        for i, sub_key in enumerate(keys):
-            if sub_key not in config:
-                raise PropertyNotFoundError(f"Key {'.'.join(keys)} not in config")
-
-            if i == len(keys) - 1:
-                return config[sub_key]
-
-            config = config[sub_key]
-
-    def add_property(self, key: str | Sequence[str], value: Any) -> None:
-        with self.secure() as toml:
-            config: dict[str, Any] = toml
+        config = self.file.read() if self.file.exists() else Document()
+        try:
+            return config.entry(key)
+        except KeyError:
             keys = split_key(key)
+            raise PropertyNotFoundError(f"Key {'.'.join(keys)} not in config")
 
-            for i, sub_key in enumerate(keys):
-                if sub_key not in config and i < len(keys) - 1:
-                    config[sub_key] = table()
-
-                if i == len(keys) - 1:
-                    config[sub_key] = value
-                    break
-
-                config = config[sub_key]
+    def add_property(self, key: str | Sequence[str], value: Any) -> None:
+        with self.secure() as config:
+            config.install(key, value)
 
     def remove_property(self, key: str | Sequence[str]) -> None:
-        with self.secure() as toml:
-            config: dict[str, Any] = toml
+        with self.secure() as config:
             keys = split_key(key)
 
             # Descend to the leaf, recording the (parent, key) at each step.
@@ -84,13 +63,13 @@ def remove_property(self, key: str | Sequence[str]) -> None:
                     break
 
     @contextmanager
-    def secure(self) -> Iterator[TOMLDocument]:
+    def secure(self) -> Iterator[Document]:
         if self.file.exists():
             initial_config = self.file.read()
             config = self.file.read()
         else:
-            initial_config = document()
-            config = document()
+            initial_config = Document()
+            config = Document()
 
         new_file = not self.file.exists()
 

src/poetry/config/file_config_source.py

@@ -8,7 +8,7 @@
 
 
 if TYPE_CHECKING:
-    from tomlkit.items import Table
+    from tomlrt import Table
 
 
 @dataclasses.dataclass(order=True, eq=True)
@@ -34,11 +34,6 @@ def to_dict(self) -> dict[str, str | bool]:
         )
 
     def to_toml_table(self) -> Table:
-        from tomlkit import nl
-        from tomlkit import table
-
-        source_table: Table = table()
-        for key, value in self.to_dict().items():
-            source_table.add(key, value)
-        source_table.add(nl())
-        return source_table
+        from tomlrt import Table
+
+        return Table.section(self.to_dict())

src/poetry/config/source.py

@@ -12,7 +12,6 @@
 from packaging.utils import canonicalize_name
 from poetry.core.packages.dependency import Dependency
 from poetry.core.packages.dependency_group import MAIN_GROUP
-from tomlkit.toml_document import TOMLDocument
 
 from poetry.console.commands.init import InitCommand
 from poetry.console.commands.installer_command import InstallerCommand
@@ -126,10 +125,8 @@ class AddCommand(InstallerCommand, InitCommand):
 
     def handle(self) -> int:
         from poetry.core.constraints.version import parse_constraint
-        from tomlkit import array
-        from tomlkit import inline_table
-        from tomlkit import nl
-        from tomlkit import table
+        from tomlrt import Array
+        from tomlrt import Table
 
         from poetry.factory import Factory
 
@@ -148,11 +145,9 @@ def handle(self) -> int:
         if optional and group != MAIN_GROUP:
             raise ValueError("You can only add optional dependencies to the main group")
 
-        # tomlkit types are awkward to work with, treat content as a mostly untyped
-        # dictionary.
-        content: dict[str, Any] = self.poetry.file.read()
-        project_content = content.get("project", table())
-        poetry_content = content.get("tool", {}).get("poetry", table())
+        content = self.poetry.file.read()
+        project_content = content.get("project", Table.section())
+        poetry_content = content.get("tool", {}).get("poetry", Table.section())
         groups_content = content.get("dependency-groups", {})
         project_name = (
             canonicalize_name(name)
@@ -174,25 +169,24 @@ def handle(self) -> int:
                 if optional:
                     project_section = project_content.get(
                         "optional-dependencies", {}
-                    ).get(optional, array())
+                    ).get(optional, Array())
                 else:
-                    project_section = project_content.get("dependencies", array())
+                    project_section = project_content.get("dependencies", Array())
                 project_dependency_names = [
                     Dependency.create_from_pep_508(dep).name for dep in project_section
                 ]
             else:
-                project_section = array()
+                project_section = Array()
 
-            poetry_section = poetry_content.get("dependencies", table())
+            poetry_section = poetry_content.get("dependencies", Table.section())
 
         # Dependency Groups
         else:
             if groups_content or "group" not in poetry_content:
                 use_groups_section = True
                 if not groups_content:
-                    groups_content = table(is_super_table=True)
-                if group not in groups_content:
-                    groups_content[group] = array("[\n]")
+                    groups_content = Table.section()
+                groups_content.setdefault(group, Array(multiline=True))
 
                 project_dependency_names = [
                     Dependency.create_from_pep_508(dep).name
@@ -207,7 +201,7 @@ def handle(self) -> int:
             poetry_section = (
                 poetry_content.get("group", {})
                 .get(group, {})
-                .get("dependencies", table())
+                .get("dependencies", Table.section())
             )
             project_section = []
 
@@ -244,7 +238,7 @@ def handle(self) -> int:
                 assert isinstance(version, str)
                 parse_constraint(version)
 
-            constraint: dict[str, Any] = inline_table()
+            constraint: dict[str, Any] = Table.inline()
             for key, value in _constraint.items():
                 if key == "name":
                     continue
@@ -331,7 +325,7 @@ def handle(self) -> int:
 
                 # create a second constraint for tool.poetry.dependencies with keys
                 # that cannot be stored in the project section
-                poetry_constraint: dict[str, Any] = inline_table()
+                poetry_constraint: dict[str, Any] = Table.inline()
                 if not isinstance(constraint, str):
                     for key in ["allow-prereleases", "develop", "source"]:
                         if value := constraint.get(key):
@@ -363,38 +357,23 @@ def handle(self) -> int:
         if project_section:
             assert group == MAIN_GROUP
             if optional:
-                if "optional-dependencies" not in project_content:
-                    project_content["optional-dependencies"] = table()
-                if optional not in project_content["optional-dependencies"]:
-                    project_content["optional-dependencies"][optional] = project_section
+                opt_deps_table = project_content.ensure_table("optional-dependencies")
+                opt_deps_table.setdefault(optional, project_section)
             elif "dependencies" not in project_content:
                 project_content["dependencies"] = project_section
 
         if poetry_section:
-            if "tool" not in content:
-                content["tool"] = table()
-            if "poetry" not in content["tool"]:
-                content["tool"]["poetry"] = poetry_content
+            tool_table = content.ensure_table("tool")
+            tool_table.setdefault("poetry", poetry_content)
+
             if group == MAIN_GROUP:
-                if "dependencies" not in poetry_content:
-                    poetry_content["dependencies"] = poetry_section
+                poetry_content.setdefault("dependencies", poetry_section)
             else:
-                if "group" not in poetry_content:
-                    poetry_content["group"] = table(is_super_table=True)
-
-                groups = poetry_content["group"]
-
-                if group not in groups:
-                    groups[group] = table()
-                    groups.add(nl())
-
-                if "dependencies" not in groups[group]:
-                    groups[group]["dependencies"] = poetry_section
+                group_table = poetry_content.ensure_table(("group", group))
+                group_table.setdefault("dependencies", poetry_section)
 
         if groups_content and group != MAIN_GROUP:
-            if "dependency-groups" not in content:
-                content["dependency-groups"] = table()
-            content["dependency-groups"][group] = groups_content[group]
+            content.install(("dependency-groups", group), groups_content[group])
 
         self.poetry.locker.set_pyproject_data(content)
         self.installer.set_locker(self.poetry.locker)
@@ -413,7 +392,6 @@ def handle(self) -> int:
         status = self.installer.run()
 
         if status == 0 and not self.option("dry-run"):
-            assert isinstance(content, TOMLDocument)
             self.poetry.file.write(content)
 
         return status

src/poetry/console/commands/add.py

@@ -11,7 +11,7 @@
 
 from cleo.helpers import option
 from packaging.utils import canonicalize_name
-from tomlkit import inline_table
+from tomlrt import Table
 
 from poetry.console.commands.command import Command
 from poetry.console.commands.env_command import EnvCommand
@@ -23,7 +23,6 @@
     from cleo.io.inputs.option import Option
     from packaging.utils import NormalizedName
     from poetry.core.packages.package import Package
-    from tomlkit.items import InlineTable
 
     from poetry.repositories import RepositoryPool
 
@@ -253,13 +252,12 @@ def _init_pyproject(
             layout_.create(project_path, with_pyproject=False)
 
         content = layout_.generate_project_content(project_path)
-        for section, item in content.items():
-            pyproject.data.append(section, item)
+        pyproject.data.update(content)
 
         if is_interactive:
             self.line("<info>Generated file</info>")
             self.line("")
-            self.line(pyproject.data.as_string().replace("\r\n", "\n"))
+            self.line(pyproject.data.render().replace("\r\n", "\n"))
             self.line("")
 
         if is_interactive and not self.confirm("Do you confirm generation?", True):
@@ -483,13 +481,11 @@ def _format_requirements(self, requirements: list[dict[str, str]]) -> Requiremen
         requires: Requirements = {}
         for requirement in requirements:
             name = requirement.pop("name")
-            constraint: str | InlineTable
+            constraint: str | Table
             if "version" in requirement and len(requirement) == 1:
                 constraint = requirement["version"]
             else:
-                constraint = inline_table()
-                constraint.trivia.trail = "\n"
-                constraint.update(requirement)
+                constraint = Table.inline(requirement)
 
             requires[name] = constraint
 

src/poetry/console/commands/init.py

@@ -9,7 +9,6 @@
 from packaging.utils import canonicalize_name
 from poetry.core.packages.dependency import Dependency
 from poetry.core.packages.dependency_group import MAIN_GROUP
-from tomlkit.toml_document import TOMLDocument
 
 from poetry.console.commands.installer_command import InstallerCommand
 
@@ -63,7 +62,7 @@ def handle(self) -> int:
         else:
             group = self.option("group", self.default_group)
 
-        content: dict[str, Any] = self.poetry.file.read()
+        content = self.poetry.file.read()
         project_content = content.get("project", {})
         groups_content = content.get("dependency-groups", {})
         poetry_content = content.get("tool", {}).get("poetry", {})
@@ -179,7 +178,6 @@ def handle(self) -> int:
         status = self.installer.run()
 
         if not self.option("dry-run") and status == 0:
-            assert isinstance(content, TOMLDocument)
             self.poetry.file.write(content)
 
         return status

src/poetry/console/commands/remove.py

@@ -87,10 +87,10 @@ def generate_system_pyproject(self) -> None:
         package.python_versions = ".".join(str(v) for v in self.env.version_info[:3])
 
         content = Factory.create_legacy_pyproject_from_package(package=package)
-        content["tool"]["poetry"]["package-mode"] = False  # type: ignore[index]
+        content["tool"]["poetry"]["package-mode"] = False
 
         for key in preserved:
-            content["tool"]["poetry"][key] = preserved[key]  # type: ignore[index]
+            content["tool"]["poetry"][key] = preserved[key]
 
         if preserved_groups:
             content["dependency-groups"] = preserved_groups

src/poetry/console/commands/self/self_command.py

@@ -1,14 +1,12 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
-from typing import Any
 from typing import ClassVar
 
 from cleo.helpers import argument
 from cleo.helpers import option
 from cleo.io.null_io import NullIO
-from tomlkit import table
-from tomlkit.items import AoT
+from tomlrt import AoT
 
 from poetry.config.source import Source
 from poetry.console.commands.command import Command
@@ -106,13 +104,8 @@ def handle(self) -> int:
             )
             return 1
 
-        # tomlkit types are awkward to work with, treat content as a mostly untyped
-        # dictionary.
-        content: dict[str, Any] = self.poetry.pyproject.data
-        if "tool" not in content:
-            content["tool"] = table()
-        if "poetry" not in content["tool"]:
-            content["tool"]["poetry"] = table()
+        content = self.poetry.pyproject.data
+        content.ensure_table("tool.poetry")
         self.poetry.pyproject.poetry_config["source"] = sources
         self.poetry.pyproject.save()
 

src/poetry/console/commands/source/add.py

@@ -4,7 +4,7 @@
 from typing import ClassVar
 
 from cleo.helpers import argument
-from tomlkit.items import AoT
+from tomlrt import AoT
 
 from poetry.console.commands.command import Command