Refactor wheel_resolver to download from non-PyPI URLs

tools/wheel_resolver/__init__.py

@@ -71,7 +71,7 @@
 )
 @click_log.simple_verbosity_option(_LOGGER)
 def main(
-    url: typing.List[str],
+    url: typing.Tuple[str],
     package_name: str,
     package_version: typing.Optional[str],
     interpreter: typing.Tuple[str, ...],
@@ -85,39 +85,36 @@ def main(
     PyPI for PACKAGE with VERSION.
 
     """
-    for u in url:
-        response = requests.head(u)
-        if response.status_code != requests.codes.ok:
-            _LOGGER.warning(
-                "%s-%s is not available, tried %r", package_name, package_version, u
-            )
-        else:
-            click.echo(u)
-            return
+    try:
+        download_output = output.get()
+    except output.OutputNotSetError:
+        _LOGGER.error("Could not get $OUTS")
+        sys.exit(1)
 
-    # We're currently hardcoding PyPI but we should consider allowing other
-    # repositories
-    # TODO (tm-jdelapuente): allow downloads from other package repositories
     locator = distlib.locators.SimpleScrapingLocator(url="https://pypi.org/simple")
     locator.wheel_tags = list(itertools.product(interpreter, abi, platform))
-    u = wheel.url(
-        package_name=package_name,
-        package_version=package_version,
-        tags=[
-            str(x)
-            for i in interpreter
-            for x in tags.generic_tags(
-                interpreter=i,
-                abis=set(abi),
-                platforms=set(platform).union({"any"}),
-            )
-        ],
-        locator=locator,
-        prereleases=prereleases,
-    )
-
-    if not output.try_download(u):
-        _LOGGER.error("Could not download from %r", u)
-        sys.exit(1)
+    try:
+        u = wheel.url(
+            package_name=package_name,
+            package_version=package_version,
+            tags=[
+                str(x)
+                for i in interpreter
+                for x in tags.generic_tags(
+                    interpreter=i,
+                    abis=set(abi),
+                    platforms=set(platform).union({"any"}),
+                )
+            ],
+            locator=locator,
+            prereleases=prereleases,
+        )
+        url += (u,)
+        print(url)
+    except Exception as error:
+        _LOGGER.warning(error)
 
-    click.echo(u)
+    origin_url = output.download(package_name, package_version, url, download_output)
+    if origin_url:
+        click.echo(f"Downloaded {package_name}-{package_version} from {origin_url}")
+        return

tools/wheel_resolver/output.py

@@ -1,27 +1,33 @@
 import os
-import sys
 import urllib.request
 import logging
+import typing
 
 _LOGGER = logging.getLogger(__name__)
 
+
 class OutputNotSetError(RuntimeError):
     pass
 
-def try_download(url):
-    """
-    Try to download url to $OUTS. Returns false if
-    it failed.
-    """
+
+def get() -> str:
     output = os.environ.get("OUTS")
     if output is None:
         raise OutputNotSetError()
-
-    try:
-        urllib.request.urlretrieve(url, output)
-    except urllib.error.HTTPError:
-        return False
-
-    return True
+    return output
 
 
+def download(
+    package_name: str, package_version: str, url: typing.List[str], download_output: str
+) -> typing.Optional[str]:
+    """Download url to $OUTS."""
+    for u in url:
+        try:
+            urllib.request.urlretrieve(u, download_output)
+        except urllib.error.HTTPError:
+            _LOGGER.warning(
+                "%s-%s is not available in %s", package_name, package_version, u
+            )
+        else:
+            return u
+    return None