[internal] Refactor lockfile_metadata.py to be more class-based (#1…

…2611)

Prework for #12610. Using a class gives better namespacing and easier access to the `LockfileMetadata` fields via `self`.

[ci skip-rust]
[ci skip-build-wheels]

src/python/pants/backend/experimental/python/lockfile.py

@@ -10,8 +10,8 @@
 from typing import Iterable
 
 from pants.backend.experimental.python.lockfile_metadata import (
+    LockfileMetadata,
     calculate_invalidation_digest,
-    lockfile_content_with_header,
 )
 from pants.backend.python.subsystems.python_tool_base import PythonToolRequirementsBase
 from pants.backend.python.target_types import EntryPoint, PythonRequirementsField
@@ -163,16 +163,20 @@ async def generate_lockfile(
         ),
     )
 
-    lockfile_digest_contents = await Get(DigestContents, Digest, poetry_export_result.output_digest)
-    lockfile_with_header = lockfile_content_with_header(
-        python_setup.lockfile_custom_regeneration_command or req.regenerate_command,
-        req.requirements_hex_digest,
-        req.interpreter_constraints,
-        lockfile_digest_contents[0].content,
+    initial_lockfile_digest_contents = await Get(
+        DigestContents, Digest, poetry_export_result.output_digest
     )
-    final_lockfile = await Get(Digest, CreateDigest([FileContent(req.dest, lockfile_with_header)]))
-
-    return PythonLockfile(final_lockfile, req.dest)
+    metadata = LockfileMetadata(req.requirements_hex_digest, req.interpreter_constraints)
+    lockfile_with_header = metadata.add_header_to_lockfile(
+        initial_lockfile_digest_contents[0].content,
+        regenerate_command=(
+            python_setup.lockfile_custom_regeneration_command or req.regenerate_command
+        ),
+    )
+    final_lockfile_digest = await Get(
+        Digest, CreateDigest([FileContent(req.dest, lockfile_with_header)])
+    )
+    return PythonLockfile(final_lockfile_digest, req.dest)
 
 
 # --------------------------------------------------------------------------------------

src/python/pants/backend/experimental/python/lockfile_metadata.py

@@ -9,7 +9,6 @@
 from typing import Any, Callable, Iterable, TypeVar
 
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
-from pants.util.ordered_set import FrozenOrderedSet
 
 BEGIN_LOCKFILE_HEADER = b"# --- BEGIN PANTS LOCKFILE METADATA: DO NOT EDIT OR REMOVE ---"
 END_LOCKFILE_HEADER = b"# --- END PANTS LOCKFILE METADATA ---"
@@ -20,15 +19,24 @@ class LockfileMetadata:
     requirements_invalidation_digest: str | None
     valid_for_interpreter_constraints: InterpreterConstraints | None
 
-    @staticmethod
-    def from_json_bytes(json_literal: str | bytes) -> LockfileMetadata:
-        """Reads a `string`/`bytes` that has JSON-encoded a `LockfileMetadata` object, returning
-        that deserialized object."""
+    @classmethod
+    def from_lockfile(cls, lockfile: bytes) -> LockfileMetadata:
+        """Parse all relevant metadata from the lockfile's header."""
+        in_metadata_block = False
+        metadata_lines = []
+        for line in lockfile.splitlines():
+            if line == BEGIN_LOCKFILE_HEADER:
+                in_metadata_block = True
+            elif line == END_LOCKFILE_HEADER:
+                break
+            elif in_metadata_block:
+                metadata_lines.append(line[2:])
+
         try:
-            metadata = json.loads(json_literal)
+            metadata = json.loads(b"\n".join(metadata_lines))
         except json.decoder.JSONDecodeError:
-            # If this block is invalid, this should trigger error/warning behavior
-            metadata = {}
+            # TODO(#12314): Add a good error.
+            raise
 
         T = TypeVar("T")
 
@@ -48,7 +56,20 @@ def coerce(t: Callable[[Any], T], key: str) -> T | None:
             ),
         )
 
-    def to_json_literal(self) -> str:
+    def add_header_to_lockfile(self, lockfile: bytes, *, regenerate_command: str) -> bytes:
+        metadata_as_a_comment = "\n".join(f"# {i}" for i in self._to_json().splitlines()).encode(
+            "ascii"
+        )
+        header = b"%b\n%b\n%b" % (BEGIN_LOCKFILE_HEADER, metadata_as_a_comment, END_LOCKFILE_HEADER)
+
+        regenerate_command_bytes = (
+            f"# This lockfile was autogenerated by Pants. To regenerate, run:\n#\n"
+            f"#    {regenerate_command}"
+        ).encode("utf-8")
+
+        return b"%b\n#\n%b\n\n%b" % (regenerate_command_bytes, header, lockfile)
+
+    def _to_json(self) -> str:
         """Produces a JSON-encoded dictionary that represents the contents of this metadata."""
         constraints = self.valid_for_interpreter_constraints
         metadata = {
@@ -57,11 +78,7 @@ def to_json_literal(self) -> str:
             if constraints
             else None,
         }
-        return json.dumps(
-            metadata,
-            ensure_ascii=True,
-            indent=2,
-        )
+        return json.dumps(metadata, ensure_ascii=True, indent=2)
 
     def is_valid_for(
         self,
@@ -92,75 +109,11 @@ def is_valid_for(
         )
 
 
-def calculate_invalidation_digest(
-    requirements: FrozenOrderedSet[str],
-) -> str:
+def calculate_invalidation_digest(requirements: Iterable[str]) -> str:
     """Returns an invalidation digest for the given requirements."""
     m = hashlib.sha256()
     inputs = {
         "requirements": list(requirements),
     }
     m.update(json.dumps(inputs).encode("utf-8"))
     return m.hexdigest()
-
-
-def lockfile_content_with_header(
-    regenerate_command: str,
-    invalidation_digest: str,
-    interpreter_constraints: InterpreterConstraints,
-    content: bytes,
-) -> bytes:
-    """Returns a version of the lockfile with Pants metadata and usage instructions prepended."""
-    regenerate_command = (
-        f"# This lockfile was autogenerated by Pants. To regenerate, run:\n#\n"
-        f"#    {regenerate_command}"
-    )
-    return b"%b\n#\n%b\n\n%b" % (
-        regenerate_command.encode("utf-8"),
-        lockfile_metadata_header(LockfileMetadata(invalidation_digest, interpreter_constraints)),
-        content,
-    )
-
-
-def lockfile_metadata_header(metadata: LockfileMetadata) -> bytes:
-    """Produces a metadata bytes object for including at the top of a lockfile.
-
-    Currently, this only consists of an invalidation digest for the file, which is used when Pants
-    consumes the lockfile during builds.
-    """
-
-    metadata_as_comment = "\n".join(f"# {i}" for i in metadata.to_json_literal().splitlines())
-
-    return (
-        b"""
-%(BEGIN_LOCKFILE_HEADER)b
-%(metadata_as_comment)s
-%(END_LOCKFILE_HEADER)b
-    """
-        % {
-            b"BEGIN_LOCKFILE_HEADER": BEGIN_LOCKFILE_HEADER,
-            b"metadata_as_comment": metadata_as_comment.encode("ascii"),
-            b"END_LOCKFILE_HEADER": END_LOCKFILE_HEADER,
-        }
-    ).strip()
-
-
-def read_lockfile_metadata(contents: bytes) -> LockfileMetadata:
-    """Reads through `contents`, and returns the contents of the lockfile metadata block as a
-    `LockfileMetadata` object."""
-
-    def yield_metadata_lines() -> Iterable[bytes]:
-        """Splits contents into lines and yields only the lines that are strictly inside the
-        metadata header block."""
-        in_metadata_block = False
-        for line in contents.splitlines():
-            if line == BEGIN_LOCKFILE_HEADER:
-                in_metadata_block = True
-            elif line == END_LOCKFILE_HEADER:
-                break
-            elif in_metadata_block:
-                yield line[2:]
-
-    metadata_lines = b"\n".join(yield_metadata_lines())
-
-    return LockfileMetadata.from_json_bytes(metadata_lines)

src/python/pants/backend/experimental/python/lockfile_metadata_test.py

@@ -1,14 +1,13 @@
 # Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
+from __future__ import annotations
+
 import pytest
 
 from pants.backend.experimental.python.lockfile_metadata import (
     LockfileMetadata,
     calculate_invalidation_digest,
-    lockfile_content_with_header,
-    lockfile_metadata_header,
-    read_lockfile_metadata,
 )
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
 from pants.util.ordered_set import FrozenOrderedSet
@@ -19,18 +18,19 @@ def test_metadata_header_round_trip() -> None:
         "cab0c0c0c0c0dadacafec0c0c0c0cafedadabeefc0c0c0c0feedbeeffeedbeef",
         InterpreterConstraints(["CPython==2.7.*", "PyPy", "CPython>=3.6,<4,!=3.7.*"]),
     )
-    serialized_metadata = lockfile_metadata_header(input_metadata)
-    output_metadata = read_lockfile_metadata(serialized_metadata)
-
+    serialized_lockfile = input_metadata.add_header_to_lockfile(
+        b"req1==1.0", regenerate_command="./pants lock"
+    )
+    output_metadata = LockfileMetadata.from_lockfile(serialized_lockfile)
     assert input_metadata == output_metadata
 
 
-def test_validated_lockfile_content() -> None:
-    content = b"""dave==3.1.4 \\
+def test_add_header_to_lockfile() -> None:
+    input_lockfile = b"""dave==3.1.4 \\
     --hash=sha256:cab0c0c0c0c0dadacafec0c0c0c0cafedadabeefc0c0c0c0feedbeeffeedbeef \\
     """
 
-    output = b"""
+    expected = b"""
 # This lockfile was autogenerated by Pants. To regenerate, run:
 #
 #    ./pants lock
@@ -47,13 +47,12 @@ def test_validated_lockfile_content() -> None:
     --hash=sha256:cab0c0c0c0c0dadacafec0c0c0c0cafedadabeefc0c0c0c0feedbeeffeedbeef \\
     """
 
-    # Helper function to make the test case more resilient to reformatting
-    line_by_line = lambda b: [i for i in (j.strip() for j in b.splitlines()) if i]
-    assert line_by_line(
-        lockfile_content_with_header(
-            "./pants lock", "000faaafcacacaca", InterpreterConstraints([">=3.7"]), content
-        )
-    ) == line_by_line(output)
+    def line_by_line(b: bytes) -> list[bytes]:
+        return [i for i in (j.strip() for j in b.splitlines()) if i]
+
+    metadata = LockfileMetadata("000faaafcacacaca", InterpreterConstraints([">=3.7"]))
+    result = metadata.add_header_to_lockfile(input_lockfile, regenerate_command="./pants lock")
+    assert line_by_line(result) == line_by_line(expected)
 
 
 _requirements = ["flake8-pantsbuild>=2.0,<3", "flake8-2020>=1.6.0,<1.7.0"]

src/python/pants/backend/python/subsystems/python_tool_base.py

@@ -13,7 +13,6 @@
 from pants.engine.fs import FileContent
 from pants.option.errors import OptionsError
 from pants.option.subsystem import Subsystem
-from pants.util.ordered_set import FrozenOrderedSet
 
 
 class PythonToolRequirementsBase(Subsystem):
@@ -135,7 +134,7 @@ def pex_requirements(
         if not self.uses_lockfile:
             return PexRequirements(requirements)
 
-        hex_digest = calculate_invalidation_digest(FrozenOrderedSet(requirements))
+        hex_digest = calculate_invalidation_digest(requirements)
 
         if self.lockfile == "<default>":
             assert self.default_lockfile_resource is not None

src/python/pants/backend/python/util_rules/pex.py

@@ -17,7 +17,7 @@
 import packaging.version
 from pkg_resources import Requirement
 
-from pants.backend.experimental.python.lockfile_metadata import read_lockfile_metadata
+from pants.backend.experimental.python.lockfile_metadata import LockfileMetadata
 from pants.backend.python.target_types import MainSpecification
 from pants.backend.python.target_types import PexPlatformsField as PythonPlatformsField
 from pants.backend.python.target_types import PythonRequirementsField
@@ -434,9 +434,8 @@ async def build_pex(
             description_of_origin=request.requirements.file_path_description_of_origin,
         )
 
-        # use contents to invalidate here
         requirements_file_digest_contents = await Get(DigestContents, PathGlobs, globs)
-        metadata = read_lockfile_metadata(requirements_file_digest_contents[0].content)
+        metadata = LockfileMetadata.from_lockfile(requirements_file_digest_contents[0].content)
         if not metadata.is_valid_for(
             request.requirements.lockfile_hex_digest,
             request.interpreter_constraints,
@@ -450,10 +449,10 @@ async def build_pex(
         requirements_file_digest = await Get(Digest, PathGlobs, globs)
 
     elif request.requirements.file_content:
-        content = request.requirements.file_content
-        argv.extend(["--requirement", content.path])
+        file_content = request.requirements.file_content
+        argv.extend(["--requirement", file_content.path])
 
-        metadata = read_lockfile_metadata(content.content)
+        metadata = LockfileMetadata.from_lockfile(file_content.content)
         if not metadata.is_valid_for(
             request.requirements.lockfile_hex_digest,
             request.interpreter_constraints,
@@ -464,7 +463,7 @@ async def build_pex(
             elif python_setup.invalid_lockfile_behavior == InvalidLockfileBehavior.warn:
                 logger.warning("%s", "Invalid lockfile provided. [TODO(#12314): Improve message]")
 
-        requirements_file_digest = await Get(Digest, CreateDigest([content]))
+        requirements_file_digest = await Get(Digest, CreateDigest([file_content]))
     else:
         argv.extend(request.requirements.req_strings)