Repository: pantsbuild/pants

Issue:
Refactor `LockfileMetadata` class hierarchy to make it easier to implement new metadata versions correctly
Currently, the `_header_dict()` method on `LockfileMetadata` needs to be called at all levels of the class hierarchy in order to function correctly. This results in a lot of guardrails being necessary, and we can probably avoid that.


Discussion:

Resolving patch:
diff --git a/src/python/pants/backend/python/util_rules/lockfile_metadata.py b/src/python/pants/backend/python/util_rules/lockfile_metadata.py
--- a/src/python/pants/backend/python/util_rules/lockfile_metadata.py
+++ b/src/python/pants/backend/python/util_rules/lockfile_metadata.py
@@ -59,16 +59,13 @@ def from_lockfile(
             ),
         )
 
-    def _header_dict(self) -> dict[Any, Any]:
-        """Produce a dictionary to be serialized into the lockfile header.
-
-        Subclasses should call `super` and update the resulting dictionary.
-        """
+    @classmethod
+    def additional_header_attrs(cls, instance: LockfileMetadata) -> dict[Any, Any]:
+        instance = cast(PythonLockfileMetadata, instance)
         return {
-            **super()._header_dict(),
             "valid_for_interpreter_constraints": [
-                str(ic) for ic in self.valid_for_interpreter_constraints
-            ],
+                str(ic) for ic in instance.valid_for_interpreter_constraints
+            ]
         }
 
     def is_valid_for(
@@ -106,11 +103,10 @@ def _from_json_dict(
 
         return PythonLockfileMetadataV1(interpreter_constraints, requirements_digest)
 
-    def _header_dict(self) -> dict[Any, Any]:
-        return {
-            **super()._header_dict(),
-            "requirements_invalidation_digest": self.requirements_invalidation_digest,
-        }
+    @classmethod
+    def additional_header_attrs(cls, instance: LockfileMetadata) -> dict[Any, Any]:
+        instance = cast(PythonLockfileMetadataV1, instance)
+        return {"requirements_invalidation_digest": instance.requirements_invalidation_digest}
 
     def is_valid_for(
         self,
@@ -166,14 +162,17 @@ def _from_json_dict(
 
         return PythonLockfileMetadataV2(interpreter_constraints, requirements)
 
-    def _header_dict(self) -> dict[Any, Any]:
+    @classmethod
+    def additional_header_attrs(cls, instance: LockfileMetadata) -> dict[Any, Any]:
+        instance = cast(PythonLockfileMetadataV2, instance)
+        # Requirements need to be stringified then sorted so that tests are deterministic. Sorting
+        # followed by stringifying does not produce a meaningful result.
         return {
-            **super()._header_dict(),
-            # Requirements need to be stringified then sorted so that tests are deterministic.
-            # Sorting followed by stringifying does not produce a meaningful result.
             "generated_with_requirements": (
-                sorted(str(i) for i in self.requirements) if self.requirements is not None else None
-            ),
+                sorted(str(i) for i in instance.requirements)
+                if instance.requirements is not None
+                else None
+            )
         }
 
     def is_valid_for(
diff --git a/src/python/pants/core/util_rules/lockfile_metadata.py b/src/python/pants/core/util_rules/lockfile_metadata.py
--- a/src/python/pants/core/util_rules/lockfile_metadata.py
+++ b/src/python/pants/core/util_rules/lockfile_metadata.py
@@ -7,7 +7,7 @@
 import json
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Callable, ClassVar, Generic, Iterable, Tuple, Type, TypeVar
+from typing import Any, Callable, ClassVar, Generic, Iterable, Tuple, Type, TypeVar, cast
 
 from pants.util.ordered_set import FrozenOrderedSet
 
@@ -151,7 +151,7 @@ def _from_json_dict(
         )
 
     def add_header_to_lockfile(self, lockfile: bytes, *, regenerate_command: str) -> bytes:
-        metadata_dict = self._header_dict()
+        metadata_dict = self.__render_header_dict()
         metadata_json = json.dumps(metadata_dict, ensure_ascii=True, indent=2).splitlines()
         metadata_as_a_comment = "\n".join(f"# {l}" for l in metadata_json).encode("ascii")
         header = b"%b\n%b\n%b" % (BEGIN_LOCKFILE_HEADER, metadata_as_a_comment, END_LOCKFILE_HEADER)
@@ -163,23 +163,42 @@ def add_header_to_lockfile(self, lockfile: bytes, *, regenerate_command: str) ->
 
         return b"%b\n#\n%b\n\n%b" % (regenerate_command_bytes, header, lockfile)
 
-    def _header_dict(self) -> dict[Any, Any]:
+    def __render_header_dict(self) -> dict[Any, Any]:
         """Produce a dictionary to be serialized into the lockfile header.
 
-        Subclasses should call `super` and update the resulting dictionary.
+        Each class should implement a class method called `additional_header_attrs`, which returns a
+        `dict` containing the metadata attributes that should be stored in the lockfile.
         """
 
-        version: int
+        attrs: dict[Any, Tuple[Any, Type]] = {}  # attr name -> (value, where we first saw it)
+        for cls in reversed(self.__class__.__mro__[:-1]):
+            new_attrs = cast(LockfileMetadata, cls).additional_header_attrs(self)
+            for attr in new_attrs:
+                if attr in attrs and attrs[attr][0] != new_attrs[attr]:
+                    raise AssertionError(
+                        f"Lockfile header attribute `{attr}` was returned by both "
+                        f"`{attrs[attr][1]}` and `{cls}`, returning different values. If these "
+                        "classes return the same attribute, they must also return the same "
+                        "value."
+                    )
+                attrs[attr] = new_attrs[attr], cls
+
+        return {key: val[0] for key, val in attrs.items()}
+
+    @classmethod
+    def additional_header_attrs(cls, instance: LockfileMetadata) -> dict[Any, Any]:
+        return {"version": instance.metadata_version()}
+
+    def metadata_version(self):
+        """Returns the version number for this metadata class, or raises an exception.
+
+        To avoid raising an exception, ensure the subclass is decorated with
+        `lockfile_metadata_version`
+        """
         for (scope, ver), cls in _concrete_metadata_classes.items():
             if isinstance(self, cls):
-                version = ver
-                break
-        else:
-            raise ValueError("Trying to serialize an unregistered `LockfileMetadata` subclass.")
-
-        return {
-            "version": version,
-        }
+                return ver
+        raise ValueError("Trying to serialize an unregistered `LockfileMetadata` subclass.")
 
 
 def calculate_invalidation_digest(requirements: Iterable[str]) -> str:
diff --git a/src/python/pants/jvm/resolve/lockfile_metadata.py b/src/python/pants/jvm/resolve/lockfile_metadata.py
--- a/src/python/pants/jvm/resolve/lockfile_metadata.py
+++ b/src/python/pants/jvm/resolve/lockfile_metadata.py
@@ -54,14 +54,6 @@ def from_lockfile(
             ),
         )
 
-    def _header_dict(self) -> dict[Any, Any]:
-        """Produce a dictionary to be serialized into the lockfile header.
-
-        Subclasses should call `super` and update the resulting dictionary.
-        """
-        d = super()._header_dict()
-        return d
-
     def is_valid_for(
         self,
         requirements: Iterable[ArtifactRequirement] | None,
@@ -106,12 +98,13 @@ def _from_json_dict(
 
         return JVMLockfileMetadataV1(requirements)
 
-    def _header_dict(self) -> dict[Any, Any]:
+    @classmethod
+    def additional_header_attrs(cls, instance: LockfileMetadata) -> dict[Any, Any]:
+        instance = cast(JVMLockfileMetadataV1, instance)
         return {
-            **super()._header_dict(),
             "generated_with_requirements": (
-                sorted(self.requirements) if self.requirements is not None else None
-            ),
+                sorted(instance.requirements) if instance.requirements is not None else None
+            )
         }
 
     def is_valid_for(


Test patch: