use tomli and tomliw to read and write lockfile

Author: dimbleby

poetry.lock

@@ -65,6 +65,8 @@ platformdirs = "^2.5.2"
 requests = "^2.18"
 requests-toolbelt = "^0.9.1"
 shellingham = "^1.5"
+tomli = "^2.0.1"
+tomli-w = "^1.0.0"
 # exclude 0.11.2 and 0.11.3 due to https://github.com/sdispater/tomlkit/issues/225
 tomlkit = ">=0.11.1,<1.0.0,!=0.11.2,!=0.11.3"
 trove-classifiers = "^2022.5.19"

pyproject.toml

@@ -233,7 +233,7 @@ def handle(self) -> int:
 
         # Refresh the locker
         self.poetry.set_locker(
-            self.poetry.locker.__class__(self.poetry.locker.lock.path, poetry_content)
+            self.poetry.locker.__class__(self.poetry.locker.lock, poetry_content)
         )
         self.installer.set_locker(self.poetry.locker)
 

src/poetry/console/commands/add.py

@@ -103,7 +103,7 @@ def handle(self) -> int:
 
         # Refresh the locker
         self.poetry.set_locker(
-            self.poetry.locker.__class__(self.poetry.locker.lock.path, poetry_content)
+            self.poetry.locker.__class__(self.poetry.locker.lock, poetry_content)
         )
         self.installer.set_locker(self.poetry.locker)
 

src/poetry/console/commands/remove.py

@@ -11,30 +11,22 @@
 from typing import Any
 from typing import cast
 
+import tomli
+import tomli_w
+
 from poetry.core.packages.dependency import Dependency
 from poetry.core.packages.package import Package
 from poetry.core.semver.helpers import parse_constraint
 from poetry.core.semver.version import Version
-from poetry.core.toml.file import TOMLFile
 from poetry.core.version.markers import parse_marker
 from poetry.core.version.requirements import InvalidRequirement
-from tomlkit import array
-from tomlkit import comment
-from tomlkit import document
-from tomlkit import inline_table
-from tomlkit import item
-from tomlkit import table
-from tomlkit.exceptions import TOMLKitError
-from tomlkit.items import Array
 
 
 if TYPE_CHECKING:
     from poetry.core.packages.directory_dependency import DirectoryDependency
     from poetry.core.packages.file_dependency import FileDependency
     from poetry.core.packages.url_dependency import URLDependency
     from poetry.core.packages.vcs_dependency import VCSDependency
-    from tomlkit.items import Table
-    from tomlkit.toml_document import TOMLDocument
 
     from poetry.repositories import Repository
 
@@ -53,17 +45,17 @@ class Locker:
     _relevant_keys = [*_legacy_keys, "group"]
 
     def __init__(self, lock: str | Path, local_config: dict[str, Any]) -> None:
-        self._lock = TOMLFile(lock)
+        self._lock = lock if isinstance(lock, Path) else Path(lock)
         self._local_config = local_config
-        self._lock_data: TOMLDocument | None = None
+        self._lock_data: dict[str, Any] | None = None
         self._content_hash = self._get_content_hash()
 
     @property
-    def lock(self) -> TOMLFile:
+    def lock(self) -> Path:
         return self._lock
 
     @property
-    def lock_data(self) -> TOMLDocument:
+    def lock_data(self) -> dict[str, Any]:
         if self._lock_data is None:
             self._lock_data = self._get_lock_data()
 
@@ -73,7 +65,7 @@ def is_locked(self) -> bool:
         """
         Checks whether the locker has been locked (lockfile found).
         """
-        if not self._lock.exists():
+        if not self.lock.exists():
             return False
 
         return "package" in self.lock_data
@@ -82,7 +74,8 @@ def is_fresh(self) -> bool:
         """
         Checks whether the lock file is still up to date with the current hash.
         """
-        lock = self._lock.read()
+        with self.lock.open("rb") as f:
+            lock = tomli.load(f)
         metadata = lock.get("metadata", {})
 
         if "content-hash" in metadata:
@@ -113,7 +106,7 @@ def locked_repository(self) -> Repository:
             source_type = source.get("type")
             url = source.get("url")
             if source_type in ["directory", "file"]:
-                url = self._lock.path.parent.joinpath(url).resolve().as_posix()
+                url = self.lock.parent.joinpath(url).resolve().as_posix()
 
             package = Package(
                 info["name"],
@@ -186,7 +179,7 @@ def locked_repository(self) -> Repository:
                         package.marker = parse_marker(split_dep[1].strip())
 
             for dep_name, constraint in info.get("dependencies", {}).items():
-                root_dir = self._lock.path.parent
+                root_dir = self.lock.parent
                 if package.source_type == "directory":
                     # root dir should be the source of the package relative to the lock
                     # path
@@ -213,29 +206,23 @@ def locked_repository(self) -> Repository:
         return packages
 
     def set_lock_data(self, root: Package, packages: list[Package]) -> bool:
-        files: dict[str, Any] = table()
+        files: dict[str, Any] = {}
         package_specs = self._lock_packages(packages)
         # Retrieving hashes
         for package in package_specs:
             if package["name"] not in files:
                 files[package["name"]] = []
 
             for f in package["files"]:
-                file_metadata = inline_table()
+                file_metadata = {}
                 for k, v in sorted(f.items()):
                     file_metadata[k] = v
 
                 files[package["name"]].append(file_metadata)
 
-            if files[package["name"]]:
-                package_files = item(files[package["name"]])
-                assert isinstance(package_files, Array)
-                files[package["name"]] = package_files.multiline(True)
-
             del package["files"]
 
-        lock = document()
-        lock.add(comment(GENERATED_COMMENT))
+        lock: dict[str, Any] = {}
         lock["package"] = package_specs
 
         if root.extras:
@@ -258,12 +245,10 @@ def set_lock_data(self, root: Package, packages: list[Package]) -> bool:
 
         return False
 
-    def _write_lock_data(self, data: TOMLDocument) -> None:
-        self.lock.write(data)
-
-        # Checking lock file data consistency
-        if data != self.lock.read():
-            raise RuntimeError("Inconsistent lock file data.")
+    def _write_lock_data(self, data: dict[str, Any]) -> None:
+        with self.lock.open("wb") as f:
+            f.write(f"# {GENERATED_COMMENT}\n\n".encode())
+            tomli_w.dump(data, f)
 
         self._lock_data = None
 
@@ -284,16 +269,17 @@ def _get_content_hash(self) -> str:
 
         return sha256(json.dumps(relevant_content, sort_keys=True).encode()).hexdigest()
 
-    def _get_lock_data(self) -> TOMLDocument:
-        if not self._lock.exists():
+    def _get_lock_data(self) -> dict[str, Any]:
+        if not self.lock.exists():
             raise RuntimeError("No lockfile found. Unable to read locked packages")
 
-        try:
-            lock_data: TOMLDocument = self._lock.read()
-        except TOMLKitError as e:
-            raise RuntimeError(f"Unable to read the lock file ({e}).")
+        with self.lock.open("rb") as f:
+            try:
+                lock_data = tomli.load(f)
+            except tomli.TOMLDecodeError as e:
+                raise RuntimeError(f"Unable to read the lock file ({e}).")
 
-        metadata = cast("Table", lock_data["metadata"])
+        metadata = lock_data["metadata"]
         lock_version = Version.parse(metadata.get("lock-version", "1.0"))
         current_version = Version.parse(self._VERSION)
         # We expect the locker to be able to read lock files
@@ -348,7 +334,7 @@ def _dump_package(self, package: Package) -> dict[str, Any]:
             if dependency.pretty_name not in dependencies:
                 dependencies[dependency.pretty_name] = []
 
-            constraint = inline_table()
+            constraint: dict[str, Any] = {}
 
             if dependency.is_directory():
                 dependency = cast("DirectoryDependency", dependency)
@@ -414,14 +400,10 @@ def _dump_package(self, package: Package) -> dict[str, Any]:
         }
 
         if dependencies:
-            data["dependencies"] = table()
-            for k, constraints in dependencies.items():
-                if len(constraints) == 1:
-                    data["dependencies"][k] = constraints[0]
-                else:
-                    data["dependencies"][k] = array().multiline(True)
-                    for constraint in constraints:
-                        data["dependencies"][k].append(constraint)
+            data["dependencies"] = {
+                name: constraints[0] if len(constraints) == 1 else constraints
+                for name, constraints in dependencies.items()
+            }
 
         if package.extras:
             extras = {}
@@ -437,7 +419,7 @@ def _dump_package(self, package: Package) -> dict[str, Any]:
                 url = Path(
                     os.path.relpath(
                         Path(url).resolve(),
-                        Path(self._lock.path.parent).resolve(),
+                        Path(self.lock.parent).resolve(),
                     )
                 ).as_posix()
 

src/poetry/packages/locker.py

@@ -416,7 +416,7 @@ def _factory(
         poetry = Factory().create_poetry(project_dir)
 
         locker = TestLocker(
-            poetry.locker.lock.path, locker_config or poetry.locker._local_config
+            poetry.locker.lock, locker_config or poetry.locker._local_config
         )
         locker.write()
 

tests/conftest.py

@@ -182,7 +182,7 @@ def reset_poetry(self) -> None:
         self._poetry.set_pool(poetry.pool)
         self._poetry.set_config(poetry.config)
         self._poetry.set_locker(
-            TestLocker(poetry.locker.lock.path, self._poetry.local_config)
+            TestLocker(poetry.locker.lock, self._poetry.local_config)
         )