add PEP 658 support!!!

Author: cosmicexplorer

src/pip/_internal/commands/download.py

@@ -1,20 +1,132 @@
 import json
 import logging
 import os
+from dataclasses import dataclass, field
 from optparse import Values
-from typing import Dict, List
+from typing import Any, Dict, List, Optional, Union
+
+from pip._vendor.packaging.requirements import Requirement
 
 from pip._internal.cli import cmdoptions
 from pip._internal.cli.cmdoptions import make_target_python
 from pip._internal.cli.req_command import RequirementCommand, with_cleanup
 from pip._internal.cli.status_codes import SUCCESS
+from pip._internal.models.link import Link
 from pip._internal.req.req_tracker import get_requirement_tracker
 from pip._internal.utils.misc import ensure_dir, normalize_path, write_output
 from pip._internal.utils.temp_dir import TempDirectory
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class RequirementHash:
+    hash_name: str
+    hash_value: str
+
+    @classmethod
+    def from_dist_info_metadata(
+        cls,
+        dist_info_metadata: Optional[str],
+    ) -> Optional["RequirementHash"]:
+        if dist_info_metadata is None:
+            return None
+        if dist_info_metadata == "true":
+            return None
+        # FIXME: don't use private `_hash_re`!
+        hash_match = Link._hash_re.match(dist_info_metadata)
+        if hash_match is None:
+            return None
+        hash_name, hash_value = hash_match.groups()
+        return cls(hash_name=hash_name, hash_value=hash_value)
+
+    @classmethod
+    def from_link(cls, link: Link) -> Optional["RequirementHash"]:
+        if not link.is_hash_allowed:
+            return None
+        hash_name = link.hash_name
+        hash_value = link.hash
+        assert hash_name is not None
+        assert hash_value is not None
+        return cls(hash_name=hash_name, hash_value=hash_value)
+
+    def as_json(self) -> Dict[str, str]:
+        return {
+            "hash_name": self.hash_name,
+            "hash_value": self.hash_value,
+        }
+
+
+@dataclass
+class DistInfoMetadata:
+    """???/From PEP 658"""
+
+    metadata_url: str
+    metadata_hash: Optional[RequirementHash]
+
+    @classmethod
+    def from_link(cls, link: Link) -> Optional["DistInfoMetadata"]:
+        if link.dist_info_metadata is None:
+            return None
+        metadata_url = f"{link.url_without_fragment}.metadata"
+        metadata_hash = RequirementHash.from_dist_info_metadata(link.dist_info_metadata)
+        return cls(metadata_url=metadata_url, metadata_hash=metadata_hash)
+
+    def as_json(self) -> Dict[str, Union[str, Optional[Dict[str, str]]]]:
+        return {
+            "metadata_url": self.metadata_url,
+            "metadata_hash": (
+                self.metadata_hash.as_json() if self.metadata_hash else None
+            ),
+        }
+
+
+@dataclass
+class RequirementDownloadInfo:
+    req: Requirement
+    url: str
+    file_hash: Optional[RequirementHash]
+    dist_info_metadata: Optional[DistInfoMetadata]
+
+    @classmethod
+    def from_req_and_link(
+        cls,
+        req: Requirement,
+        link: Link,
+    ) -> "RequirementDownloadInfo":
+        return cls(
+            req=req,
+            url=link.url,
+            file_hash=RequirementHash.from_link(link),
+            dist_info_metadata=DistInfoMetadata.from_link(link),
+        )
+
+    def as_json(self) -> Dict[str, Any]:
+        return {
+            "req": str(self.req),
+            "url": self.url,
+            "hash": self.file_hash and self.file_hash.as_json(),
+            "dist_info_metadata": (
+                self.dist_info_metadata and self.dist_info_metadata.as_json()
+            ),
+        }
+
+
+@dataclass
+class DownloadInfos:
+    # python_version: Optional[Requirement] = None
+    python_version: Optional[str] = None
+    resolution: Dict[str, RequirementDownloadInfo] = field(default_factory=dict)
+
+    def as_json(self) -> Dict[str, Any]:
+        return {
+            "python_version": self.python_version and str(self.python_version),
+            "resolution": {
+                name: info.as_json() for name, info in self.resolution.items()
+            },
+        }
+
+
 class DownloadCommand(RequirementCommand):
     """
     Download packages from:
@@ -149,24 +261,37 @@ def run(self, options: Values, args: List[str]) -> int:
         requirement_set = resolver.resolve(reqs, check_supported_wheels=True)
 
         downloaded: List[str] = []
-        download_infos: List[Dict[str, str]] = []
         for req in requirement_set.requirements.values():
+            # If this distribution was not already satisfied, that means we
+            # downloaded it.
             if req.satisfied_by is None:
-                assert req.name is not None
-                assert req.link is not None
-                download_infos.append(
-                    {
-                        "name": req.name,
-                        "url": req.link.url,
-                    }
-                )
                 preparer.save_linked_requirement(req)
+                assert req.name is not None
                 downloaded.append(req.name)
 
+        download_infos = DownloadInfos()
+        assert requirement_set.candidates is not None
+        for candidate in requirement_set.candidates.mapping.values():
+            # This will occur for the python version requirement, for example.
+            if candidate.name not in requirement_set.requirements:
+                # download_infos.python_version = candidate.as_requirement()
+                download_infos.python_version = str(candidate)
+                continue
+            req = requirement_set.requirements[candidate.name]
+            assert req.name is not None
+            assert req.link is not None
+            assert req.name not in download_infos.resolution
+            download_infos.resolution[
+                req.name
+            ] = RequirementDownloadInfo.from_req_and_link(
+                req=candidate.as_requirement(),
+                link=req.link,
+            )
+
         if downloaded:
             write_output("Successfully downloaded %s", " ".join(downloaded))
         if options.print_download_urls:
             with open(options.print_download_urls, "w") as f:
-                json.dump(download_infos, f, indent=4)
+                json.dump(download_infos.as_json(), f, indent=4)
 
         return SUCCESS

src/pip/_internal/index/collector.py

@@ -248,12 +248,14 @@ def _create_link_from_element(
     url = _clean_link(urllib.parse.urljoin(base_url, href))
     pyrequire = anchor.get("data-requires-python")
     yanked_reason = anchor.get("data-yanked")
+    dist_info_metadata = anchor.get("data-dist-info-metadata")
 
     link = Link(
         url,
         comes_from=page_url,
         requires_python=pyrequire,
         yanked_reason=yanked_reason,
+        dist_info_metadata=dist_info_metadata,
     )
 
     return link

src/pip/_internal/models/link.py

@@ -34,6 +34,7 @@ class Link(KeyBasedCompareMixin):
         "comes_from",
         "requires_python",
         "yanked_reason",
+        "dist_info_metadata",
         "cache_link_parsing",
     ]
 
@@ -43,6 +44,7 @@ def __init__(
         comes_from: Optional[Union[str, "HTMLPage"]] = None,
         requires_python: Optional[str] = None,
         yanked_reason: Optional[str] = None,
+        dist_info_metadata: Optional[str] = None,
         cache_link_parsing: bool = True,
     ) -> None:
         """
@@ -59,6 +61,7 @@ def __init__(
             a simple repository HTML link. If the file has been yanked but
             no reason was provided, this should be the empty string. See
             PEP 592 for more information and the specification.
+        :param dist_info_metadata: ???/PEP 658
         :param cache_link_parsing: A flag that is used elsewhere to determine
                                    whether resources retrieved from this link
                                    should be cached. PyPI index urls should
@@ -78,6 +81,7 @@ def __init__(
         self.comes_from = comes_from
         self.requires_python = requires_python if requires_python else None
         self.yanked_reason = yanked_reason
+        self.dist_info_metadata = dist_info_metadata
 
         super().__init__(key=url, defining_class=Link)
 
@@ -165,6 +169,7 @@ def subdirectory_fragment(self) -> Optional[str]:
             return None
         return match.group(1)
 
+    # FIXME: retrieve all the `re.compile` anchor tags when the Link is constructed!!
     _hash_re = re.compile(
         r"({choices})=([a-f0-9]+)".format(choices="|".join(_SUPPORTED_HASHES))
     )

src/pip/_internal/req/req_install.py

@@ -62,6 +62,23 @@
 logger = logging.getLogger(__name__)
 
 
+def produce_exact_version_requirement(name: str, version: str) -> Requirement:
+    if isinstance(parse_version(version), Version):
+        op = "=="
+    else:
+        op = "==="
+
+    return Requirement(
+        "".join(
+            [
+                name,
+                op,
+                version,
+            ]
+        )
+    )
+
+
 class InstallRequirement:
     """
     Represents something that may be installed later on, may have information
@@ -348,20 +365,10 @@ def _set_requirement(self) -> None:
         assert self.metadata is not None
         assert self.source_dir is not None
 
-        # Construct a Requirement object from the generated metadata
-        if isinstance(parse_version(self.metadata["Version"]), Version):
-            op = "=="
-        else:
-            op = "==="
-
-        self.req = Requirement(
-            "".join(
-                [
-                    self.metadata["Name"],
-                    op,
-                    self.metadata["Version"],
-                ]
-            )
+        # Construct a Requirement object from the generated metadata.
+        self.req = produce_exact_version_requirement(
+            self.metadata["Name"],
+            self.metadata["Version"],
         )
 
     def warn_on_mismatching_name(self) -> None:

src/pip/_internal/resolution/base.py

@@ -1,17 +1,30 @@
-from typing import Callable, List, Optional
+from typing import TYPE_CHECKING, Callable, List, Optional
 
 from pip._internal.req.req_install import InstallRequirement
 from pip._internal.req.req_set import RequirementSet
 
+if TYPE_CHECKING:
+    from pip._vendor.resolvelib.resolvers import Result as RLResult
+
+    from .resolvelib.base import Candidate, Requirement
+
+    Result = RLResult[Requirement, Candidate, str]
+
 InstallRequirementProvider = Callable[
     [str, Optional[InstallRequirement]], InstallRequirement
 ]
 
 
+class RequirementSetWithCandidates(RequirementSet):
+    def __init__(self, check_supported_wheels: bool = True) -> None:
+        super().__init__(check_supported_wheels=check_supported_wheels)
+        self.candidates: Optional[Result] = None
+
+
 class BaseResolver:
     def resolve(
         self, root_reqs: List[InstallRequirement], check_supported_wheels: bool
-    ) -> RequirementSet:
+    ) -> RequirementSetWithCandidates:
         raise NotImplementedError()
 
     def get_installation_order(

src/pip/_internal/resolution/resolvelib/base.py

@@ -1,5 +1,7 @@
+import abc
 from typing import FrozenSet, Iterable, Optional, Tuple, Union
 
+from pip._vendor.packaging.requirements import Requirement as PkgRequirement
 from pip._vendor.packaging.specifiers import SpecifierSet
 from pip._vendor.packaging.utils import NormalizedName, canonicalize_name
 from pip._vendor.packaging.version import LegacyVersion, Version
@@ -95,47 +97,52 @@ def _match_link(link: Link, candidate: "Candidate") -> bool:
     return False
 
 
-class Candidate:
-    @property
+class Candidate(metaclass=abc.ABCMeta):
+    @abc.abstractproperty
     def project_name(self) -> NormalizedName:
         """The "project name" of the candidate.
 
         This is different from ``name`` if this candidate contains extras,
         in which case ``name`` would contain the ``[...]`` part, while this
         refers to the name of the project.
         """
-        raise NotImplementedError("Override in subclass")
 
-    @property
+    @abc.abstractproperty
     def name(self) -> str:
         """The name identifying this candidate in the resolver.
 
         This is different from ``project_name`` if this candidate contains
         extras, where ``project_name`` would not contain the ``[...]`` part.
         """
-        raise NotImplementedError("Override in subclass")
 
-    @property
+    @abc.abstractproperty
     def version(self) -> CandidateVersion:
-        raise NotImplementedError("Override in subclass")
+        ...
 
-    @property
+    @abc.abstractmethod
+    def as_requirement(self) -> PkgRequirement:
+        ...
+
+    @abc.abstractproperty
     def is_installed(self) -> bool:
-        raise NotImplementedError("Override in subclass")
+        ...
 
-    @property
+    @abc.abstractproperty
     def is_editable(self) -> bool:
-        raise NotImplementedError("Override in subclass")
+        ...
 
-    @property
+    @abc.abstractproperty
     def source_link(self) -> Optional[Link]:
-        raise NotImplementedError("Override in subclass")
+        ...
 
+    @abc.abstractmethod
     def iter_dependencies(self, with_requires: bool) -> Iterable[Optional[Requirement]]:
-        raise NotImplementedError("Override in subclass")
+        ...
 
+    @abc.abstractmethod
     def get_install_requirement(self) -> Optional[InstallRequirement]:
-        raise NotImplementedError("Override in subclass")
+        ...
 
+    @abc.abstractmethod
     def format_for_error(self) -> str:
-        raise NotImplementedError("Subclass should override")
+        ...