From 5ad7ef16b8e960ef54e64970ebdc3d91dc5c6ed1 Mon Sep 17 00:00:00 2001
From: Frost Ming <me@frostming.com>
Date: Fri, 12 Jul 2024 09:50:00 +0800
Subject: [PATCH] feat: generate provenance_url.json when installing
 packages(PEP 710)

Signed-off-by: Frost Ming <me@frostming.com>
---
 src/pdm/installers/installers.py |  8 +++-----
 src/pdm/installers/manager.py    | 11 +++++++++--
 src/pdm/models/candidates.py     | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/src/pdm/installers/installers.py b/src/pdm/installers/installers.py
index 7d3563ce97..d85938a0a4 100644
--- a/src/pdm/installers/installers.py
+++ b/src/pdm/installers/installers.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 import os
 import stat
 from functools import cached_property
@@ -151,7 +150,7 @@ def _get_link_method(cache_method: str) -> LinkMethod:
 def install_wheel(
     wheel: Path,
     environment: BaseEnvironment,
-    direct_url: dict[str, Any] | None = None,
+    additional_metadata: dict[str, bytes] | None = None,
     install_links: bool = False,
     rename_pth: bool = False,
 ) -> str:
@@ -169,9 +168,8 @@ def install_wheel(
     else:
         link_method = _get_link_method(cache_method)
 
-    additional_metadata: dict[str, bytes] = {}
-    if direct_url is not None:
-        additional_metadata["direct_url.json"] = json.dumps(direct_url, indent=2).encode()
+    if additional_metadata is None:
+        additional_metadata = {}
 
     destination = InstallDestination(
         scheme_dict=environment.get_paths(dist_name),
diff --git a/src/pdm/installers/manager.py b/src/pdm/installers/manager.py
index 4a2061c458..67e6b9968b 100644
--- a/src/pdm/installers/manager.py
+++ b/src/pdm/installers/manager.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 from typing import TYPE_CHECKING
 
 from pdm import termui
@@ -29,10 +30,16 @@ def __init__(
     def install(self, candidate: Candidate) -> Distribution:
         """Install a candidate into the environment, return the distribution"""
         prepared = candidate.prepare(self.environment)
+        wheel = prepared.build()
+        additional_metadata: dict[str, bytes] = {}
+        if direct_url := prepared.direct_url():
+            additional_metadata["direct_url.json"] = json.dumps(direct_url, indent=2).encode("utf-8")
+        elif provenance_url := prepared.provenance_url():
+            additional_metadata["provenance_url.json"] = json.dumps(provenance_url, indent=2).encode("utf-8")
         dist_info = install_wheel(
-            prepared.build(),
+            wheel,
             self.environment,
-            direct_url=prepared.direct_url(),
+            additional_metadata=additional_metadata,
             install_links=self.use_install_cache and not candidate.req.editable,
             rename_pth=self.rename_pth,
         )
diff --git a/src/pdm/models/candidates.py b/src/pdm/models/candidates.py
index 9d062d6d8f..a5158ad85b 100644
--- a/src/pdm/models/candidates.py
+++ b/src/pdm/models/candidates.py
@@ -3,6 +3,7 @@
 import dataclasses
 import hashlib
 import os
+import posixpath
 import re
 import warnings
 from functools import cached_property
@@ -21,6 +22,7 @@
 from pdm.models.reporter import BaseReporter
 from pdm.models.requirements import (
     FileRequirement,
+    NamedRequirement,
     Requirement,
     VcsRequirement,
     _egg_info_re,
@@ -47,6 +49,9 @@
     from pdm.environments import BaseEnvironment
 
 
+ALLOWED_HASHES = hashlib.algorithms_guaranteed - {"shake_128", "shake_256", "sha1", "md5"}
+
+
 def _dist_info_files(whl_zip: ZipFile) -> list[str]:
     """Identify the .dist-info folder inside a wheel ZipFile."""
     res = []
@@ -335,7 +340,9 @@ def revision(self) -> str:
         )
 
     def direct_url(self) -> dict[str, Any] | None:
-        """PEP 610 direct_url.json data"""
+        """PEP 610 direct_url.json data
+        https://peps.python.org/pep-0610/
+        """
         req = self.req
         if isinstance(req, VcsRequirement):
             if req.editable:
@@ -383,6 +390,29 @@ def direct_url(self) -> dict[str, Any] | None:
         else:
             return None
 
+    def provenance_url(self) -> dict[str, Any] | None:
+        """PEP 710 provenance_url.json data
+        https://peps.python.org/pep-0710/
+        """
+        req = self.req
+        if not isinstance(req, NamedRequirement):
+            return None
+        assert self.link is not None
+        comes_from = self.link.comes_from  # e.g. https://pypi.org/simple/requests/
+        if comes_from is None:  # can't determine the index_url
+            return None
+        # FIXME: what about find-links source?
+        index_url = posixpath.dirname(comes_from.rstrip("/")) + "/"
+        return {
+            "url": self.link.url_without_fragment,
+            "index_url": index_url,
+            "archive_info": {
+                "hashes": {
+                    name: hashes[0] for name, hashes in (self.link.hash_option or {}).items() if name in ALLOWED_HASHES
+                },
+            },
+        }
+
     def build(self) -> Path:
         """Call PEP 517 build hook to build the candidate into a wheel"""
         self._obtain(allow_all=False)