Merge pull request #991 from tiran/override-resolver-url

feat: add override_download_url to resolvers

Author: LalatenduMohanty

src/fromager/resolver.py

@@ -14,7 +14,7 @@
 from collections.abc import Iterable
 from operator import attrgetter
 from platform import python_version
-from urllib.parse import quote, unquote, urljoin, urlparse
+from urllib.parse import quote, unquote, urlparse
 
 import pypi_simple
 import resolvelib
@@ -192,6 +192,8 @@ def get_project_from_pypi(
     extras: typing.Iterable[str],
     sdist_server_url: str,
     ignore_platform: bool = False,
+    *,
+    override_download_url: str | None = None,
 ) -> Candidates:
     """Return candidates created from the project name and extras."""
     found_candidates: set[str] = set()
@@ -352,14 +354,19 @@ def get_project_from_pypi(
                 ignored_candidates.add(dp.filename)
                 continue
 
+        if override_download_url is None:
+            url = dp.url
+        else:
+            url = override_download_url.format(version=version)
+
         upload_time = dp.upload_time
         if upload_time is not None:
             upload_time = upload_time.astimezone(datetime.UTC)
 
         c = Candidate(
             name=name,
             version=version,
-            url=dp.url,
+            url=url,
             extras=tuple(sorted(extras)),
             is_sdist=is_sdist,
             build_tag=build_tag,
@@ -592,7 +599,11 @@ def find_matches(
 
 
 class PyPIProvider(BaseProvider):
-    """Lookup package and versions from a simple Python index (PyPI)"""
+    """Lookup package and versions from a simple Python index (PyPI)
+
+    The ``override_download_url`` parameter supports the string template variable:
+    * version (Version object)
+    """
 
     provider_description: typing.ClassVar[str] = (
         "PyPI resolver (searching at {self.sdist_server_url})"
@@ -608,6 +619,7 @@ def __init__(
         ignore_platform: bool = False,
         *,
         use_resolver_cache: bool = True,
+        override_download_url: str | None = None,
     ):
         super().__init__(
             constraints=constraints,
@@ -618,21 +630,25 @@ def __init__(
         self.include_wheels = include_wheels
         self.sdist_server_url = sdist_server_url
         self.ignore_platform = ignore_platform
+        self.override_download_url = override_download_url
 
     @property
     def cache_key(self) -> str:
         # ignore platform parameter changes behavior of find_candidates()
+        key = self.sdist_server_url
+        if self.override_download_url is not None:
+            key = f"{key}+{self.override_download_url}"
         if self.ignore_platform:
-            return f"{self.sdist_server_url}+ignore_platform"
-        else:
-            return self.sdist_server_url
+            key = f"{key}+ignore_platform"
+        return key
 
     def find_candidates(self, identifier: str) -> Candidates:
         return get_project_from_pypi(
             identifier,
-            set(),
-            self.sdist_server_url,
-            self.ignore_platform,
+            extras=set(),
+            sdist_server_url=self.sdist_server_url,
+            ignore_platform=self.ignore_platform,
+            override_download_url=self.override_download_url,
         )
 
     def validate_candidate(
@@ -791,6 +807,12 @@ class GitHubTagProvider(GenericProvider):
     """Lookup tarball and version from GitHub git tags
 
     Assumes that upstream uses version tags `1.2.3` or `v1.2.3`.
+
+    The ``override_download_url`` parameter supports the string template variable:
+    * organization
+    * repo
+    * tagname
+    * version (Version object)
     """
 
     provider_description: typing.ClassVar[str] = (
@@ -808,6 +830,7 @@ def __init__(
         *,
         req_type: RequirementType | None = None,
         use_resolver_cache: bool = True,
+        override_download_url: str | None = None,
     ):
         super().__init__(
             constraints=constraints,
@@ -818,10 +841,14 @@ def __init__(
         )
         self.organization = organization
         self.repo = repo
+        self.override_download_url = override_download_url
 
     @property
     def cache_key(self) -> str:
-        return f"{self.organization}/{self.repo}"
+        key = f"{self.organization}/{self.repo}"
+        if self.override_download_url is not None:
+            key = f"{key}+{self.override_download_url}"
+        return key
 
     @retry_on_exception(
         exceptions=RETRYABLE_EXCEPTIONS,
@@ -852,7 +879,16 @@ def _find_tags(
                     logger.debug(f"{identifier}: match function ignores {tagname}")
                     continue
                 assert isinstance(version, Version)
-                url = entry["tarball_url"]
+
+                if self.override_download_url is None:
+                    url = entry["tarball_url"]
+                else:
+                    url = self.override_download_url.format(
+                        organization=self.organization,
+                        repo=self.repo,
+                        tagname=tagname,
+                        version=version,
+                    )
 
                 # Github tag API endpoint does not include commit date information.
                 # It would be too expensive to query every commit API endpoint.
@@ -870,7 +906,15 @@ def _find_tags(
 
 
 class GitLabTagProvider(GenericProvider):
-    """Lookup tarball and version from GitLab git tags"""
+    """Lookup tarball and version from GitLab git tags
+
+    The ``override_download_url`` parameter supports the string template variable:
+    * hostname
+    * project_path
+    * project_name (last component of project_path)
+    * tagname
+    * version (Version object)
+    """
 
     provider_description: typing.ClassVar[str] = (
         "GitLab tag resolver (project: {self.server_url}/{self.project_path})"
@@ -885,6 +929,7 @@ def __init__(
         *,
         req_type: RequirementType | None = None,
         use_resolver_cache: bool = True,
+        override_download_url: str | None = None,
     ) -> None:
         super().__init__(
             constraints=constraints,
@@ -894,6 +939,9 @@ def __init__(
             matcher=matcher,
         )
         self.server_url = server_url.rstrip("/")
+        self.server_hostname = urlparse(server_url).hostname
+        if not self.server_hostname:
+            raise ValueError(f"invalid {server_url=}")
         self.project_path = project_path.lstrip("/")
         # URL-encode the project path as required by GitLab API.
         # The safe="" parameter tells quote() to encode ALL characters,
@@ -904,10 +952,14 @@ def __init__(
         self.api_url = (
             f"{self.server_url}/api/v4/projects/{encoded_path}/repository/tags"
         )
+        self.override_download_url = override_download_url
 
     @property
     def cache_key(self) -> str:
-        return f"{self.server_url}/{self.project_path}"
+        key = f"{self.server_url}/{self.project_path}"
+        if self.override_download_url is not None:
+            key = f"{key}+{self.override_download_url}"
+        return key
 
     @retry_on_exception(
         exceptions=RETRYABLE_EXCEPTIONS,
@@ -921,6 +973,14 @@ def _find_tags(
     ) -> Iterable[Candidate]:
         nexturl: str = self.api_url
         created_at: datetime.datetime | None
+        project_name = self.project_path.split("/")[-1]
+        if self.override_download_url is None:
+            download_template = (
+                self.server_url
+                + "/{project_path}/-/archive/{tagname}/{project_name}-{tagname}.tar.gz"
+            )
+        else:
+            download_template = self.override_download_url
         while nexturl:
             resp: Response = session.get(nexturl)
             resp.raise_for_status()
@@ -932,8 +992,13 @@ def _find_tags(
                     continue
                 assert isinstance(version, Version)
 
-                archive_path: str = f"{self.project_path}/-/archive/{tagname}/{self.project_path.split('/')[-1]}-{tagname}.tar.gz"
-                url = urljoin(self.server_url, archive_path)
+                url = download_template.format(
+                    hostname=self.server_hostname,
+                    project_path=self.project_path,
+                    project_name=project_name,
+                    tagname=tagname,
+                    version=version,
+                )
 
                 # get tag creation time, fall back to commit creation time
                 created_at_str: str | None = entry.get("created_at")

tests/test_resolver.py

@@ -372,6 +372,26 @@ def test_provider_constraint_match() -> None:
         assert str(candidate.version) == "1.2.2"
 
 
+def test_pypi_provider_override_download_url() -> None:
+    with requests_mock.Mocker() as r:
+        r.get(
+            "https://pypi.org/simple/hydra-core/",
+            text=_hydra_core_simple_response,
+        )
+
+        provider = resolver.PyPIProvider(
+            override_download_url="https://server.test/hydr_core-{version}.tar.gz"
+        )
+        reporter: resolvelib.BaseReporter = resolvelib.BaseReporter()
+        rslvr = resolvelib.Resolver(provider, reporter)
+
+        result = rslvr.resolve([Requirement("hydra-core")])
+        assert "hydra-core" in result.mapping
+
+        candidate = result.mapping["hydra-core"]
+        assert candidate.url == "https://server.test/hydr_core-1.3.2.tar.gz"
+
+
 _ignore_platform_simple_response = """
 <!DOCTYPE html>
 <html>
@@ -717,6 +737,33 @@ def test_resolve_github() -> None:
         )
 
 
+def test_resolve_github_override_download_url() -> None:
+    with requests_mock.Mocker() as r:
+        r.get(
+            "https://api.github.com:443/repos/python-wheel-build/fromager",
+            text=_github_fromager_repo_response,
+        )
+        r.get(
+            "https://api.github.com:443/repos/python-wheel-build/fromager/tags",
+            text=_github_fromager_tag_response,
+        )
+
+        provider = resolver.GitHubTagProvider(
+            organization="python-wheel-build",
+            repo="fromager",
+            override_download_url="git+https://github.com/{organization}/{repo}.git@{tagname}",
+        )
+        reporter: resolvelib.BaseReporter = resolvelib.BaseReporter()
+        rslvr = resolvelib.Resolver(provider, reporter)
+
+        result = rslvr.resolve([Requirement("fromager")])
+        candidate = result.mapping["fromager"]
+        assert (
+            str(candidate.url)
+            == "git+https://github.com/python-wheel-build/fromager.git@0.9.0"
+        )
+
+
 def test_github_constraint_mismatch() -> None:
     constraint = constraints.Constraints()
     constraint.add_constraint("fromager>=1.0")
@@ -1005,6 +1052,29 @@ def test_resolve_gitlab() -> None:
         )
 
 
+def test_resolve_gitlab_override_download_url() -> None:
+    with requests_mock.Mocker() as r:
+        r.get(
+            "https://gitlab.com/api/v4/projects/mirrors%2Fgithub%2Fdecile-team%2Fsubmodlib/repository/tags",
+            text=_gitlab_submodlib_repo_response,
+        )
+
+        provider = resolver.GitLabTagProvider(
+            project_path="mirrors/github/decile-team/submodlib",
+            server_url="https://gitlab.com",
+            matcher=re.compile("v(.*)"),  # with match object
+            override_download_url="git+https://{hostname}/{project_path}.git@{tagname}",
+        )
+        reporter: resolvelib.BaseReporter = resolvelib.BaseReporter()
+        rslvr = resolvelib.Resolver(provider, reporter)
+        result = rslvr.resolve([Requirement("submodlib")])
+        candidate = result.mapping["submodlib"]
+        assert (
+            str(candidate.url)
+            == "git+https://gitlab.com/mirrors/github/decile-team/submodlib.git@v0.0.3"
+        )
+
+
 def test_gitlab_constraint_mismatch() -> None:
     constraint = constraints.Constraints()
     constraint.add_constraint("submodlib>=1.0")