Add an optional --index-url argument similar to that in pip

hashin.py

@@ -25,8 +25,10 @@
 if sys.version_info >= (3,):
     from urllib.request import urlopen
     from urllib.error import HTTPError
+    from urllib.parse import urljoin
 else:
     from urllib import urlopen
+    from urlparse import urljoin
 
     input = raw_input  # noqa
 
@@ -44,6 +46,8 @@
 
 DEFAULT_ALGORITHM = "sha256"
 
+DEFAULT_INDEX_URL = os.environ.get("INDEX_URL", "https://pypi.org/")
+
 MAX_WORKERS = None
 
 if sys.version_info >= (3, 4) and sys.version_info < (3, 5):
@@ -153,6 +157,7 @@ def run_packages(
     previous_versions=None,
     interactive=False,
     synchronous=False,
+    index_url=DEFAULT_INDEX_URL,
 ):
     assert isinstance(specs, list), type(specs)
     all_new_lines = []
@@ -161,7 +166,9 @@ def run_packages(
 
     lookup_memory = {}
     if not synchronous and len(specs) > 1:
-        pre_download_packages(lookup_memory, specs, verbose=verbose)
+        pre_download_packages(
+            lookup_memory, specs, verbose=verbose, index_url=index_url
+        )
 
     for spec in specs:
         package, version, restriction = _explode_package_spec(spec)
@@ -186,6 +193,7 @@ def run_packages(
             algorithm=algorithm,
             include_prereleases=include_prereleases,
             lookup_memory=lookup_memory,
+            index_url=index_url,
         )
         package = data["package"]
         # We need to keep this `req` instance for the sake of turning it into a string
@@ -273,14 +281,14 @@ def run_packages(
     return 0
 
 
-def pre_download_packages(memory, specs, verbose=False):
+def pre_download_packages(memory, specs, verbose=False, index_url=DEFAULT_INDEX_URL):
     futures = {}
     with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
         for spec in specs:
             package, _, _ = _explode_package_spec(spec)
             req = Requirement(package)
             futures[
-                executor.submit(get_package_data, req.name, verbose=verbose)
+                executor.submit(get_package_data, req.name, index_url, verbose=verbose)
             ] = req.name
         for future in concurrent.futures.as_completed(futures):
             content = future.result()
@@ -567,8 +575,9 @@ def filter_releases(releases, python_versions):
     return filtered
 
 
-def get_package_data(package, verbose=False):
-    url = "https://pypi.org/pypi/%s/json" % package
+def get_package_data(package, index_url, verbose=False):
+    path = "/pypi/%s/json" % package
+    url = urljoin(index_url, path)
     if verbose:
         print(url)
     content = json.loads(_download(url))
@@ -615,6 +624,7 @@ def get_package_hashes(
     verbose=False,
     include_prereleases=False,
     lookup_memory=None,
+    index_url=DEFAULT_INDEX_URL,
 ):
     """
     Gets the hashes for the given package.
@@ -642,7 +652,7 @@ def get_package_hashes(
     if lookup_memory is not None and package in lookup_memory:
         data = lookup_memory[package]
     else:
-        data = get_package_data(package, verbose)
+        data = get_package_data(package, index_url, verbose)
     if not version:
         version = get_latest_version(data, include_prereleases)
         assert version
@@ -741,6 +751,11 @@ def get_parser():
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--index-url",
+        help="alternate package index url (default {0})".format(DEFAULT_INDEX_URL),
+        default=None,
+    )
     return parser
 
 
@@ -786,6 +801,7 @@ def main():
             dry_run=args.dry_run,
             interactive=args.interactive,
             synchronous=args.synchronous,
+            index_url=args.index_url,
         )
     except PackageError as exception:
         print(str(exception), file=sys.stderr)

tests/test_arg_parse.py

@@ -16,6 +16,8 @@ def test_everything():
             "3.5",
             "-v",
             "--dry-run",
+            "--index-url",
+            "https://pypi1.someorg.net/",
         ]
     )
     expected = argparse.Namespace(
@@ -30,6 +32,7 @@ def test_everything():
         update_all=False,
         interactive=False,
         synchronous=False,
+        index_url="https://pypi1.someorg.net/",
     )
     assert args == (expected, [])
 
@@ -47,6 +50,8 @@ def test_everything_long():
             "3.5",
             "--verbose",
             "--dry-run",
+            "--index-url",
+            "https://pypi1.someorg.net/",
         ]
     )
     expected = argparse.Namespace(
@@ -61,6 +66,7 @@ def test_everything_long():
         update_all=False,
         interactive=False,
         synchronous=False,
+        index_url="https://pypi1.someorg.net/",
     )
     assert args == (expected, [])
 
@@ -79,5 +85,6 @@ def test_minimal():
         update_all=False,
         interactive=False,
         synchronous=False,
+        index_url=None,
     )
     assert args == (expected, [])

tests/test_cli.py

@@ -146,6 +146,7 @@ def mock_parse_args(*a, **k):
             update_all=False,
             interactive=False,
             synchronous=False,
+            index_url=None,
         )
 
     mock_get_parser().parse_args.side_effect = mock_parse_args
@@ -1220,6 +1221,53 @@ def mocked_get(url, **options):
         assert output.startswith("hashin==0.10")
 
 
+def test_run_with_alternate_index_url(murlopen, tmpfile):
+    def mocked_get(url, **options):
+        if url == "https://pypi.internal.net/pypi/hashin/json":
+            return _Response(
+                {
+                    "info": {"version": "0.10", "name": "hashin"},
+                    "releases": {
+                        "0.10": [
+                            {
+                                "url": "https://pypi.internal.net/packages/2.7/p/hashin/hashin-0.10-py2-none-any.whl",
+                                "digests": {"sha256": "aaaaa"},
+                            },
+                            {
+                                "url": "https://pypi.internal.net/packages/3.3/p/hashin/hashin-0.10-py3-none-any.whl",
+                                "digests": {"sha256": "bbbbb"},
+                            },
+                            {
+                                "url": "https://pypi.internal.net/packages/source/p/hashin/hashin-0.10.tar.gz",
+                                "digests": {"sha256": "ccccc"},
+                            },
+                        ]
+                    },
+                }
+            )
+
+        raise NotImplementedError(url)
+
+    murlopen.side_effect = mocked_get
+
+    with tmpfile() as filename:
+        with open(filename, "w") as f:
+            f.write("")
+
+        retcode = hashin.run(
+            "hashin",
+            filename,
+            "sha256",
+            verbose=True,
+            index_url="https://pypi.internal.net/",
+        )
+
+        assert retcode == 0
+        with open(filename) as f:
+            output = f.read()
+        assert output.startswith("hashin==0.10")
+
+
 def test_run_contained_names(murlopen, tmpfile):
     """
     This is based on https://github.com/peterbe/hashin/issues/35
@@ -2005,6 +2053,51 @@ def mocked_get(url, **options):
     assert result == expected
 
 
+def test_get_package_hashes_from_alternate_index_url(murlopen):
+    def mocked_get(url, **options):
+        if url == "https://pypi.internal.net/pypi/hashin/json":
+            return _Response(
+                {
+                    "info": {"version": "0.10", "name": "hashin"},
+                    "releases": {
+                        "0.10": [
+                            {
+                                "url": "https://pypi.internal.net/packages/2.7/p/hashin/hashin-0.10-py2-none-any.whl",
+                                "digests": {"sha256": "ddddd"},
+                            },
+                            {
+                                "url": "https://pypi.internal.net/packages/3.3/p/hashin/hashin-0.10-py3-none-any.whl",
+                                "digests": {"sha256": "eeeee"},
+                            },
+                            {
+                                "url": "https://pypi.internal.net/packages/source/p/hashin/hashin-0.10.tar.gz",
+                                "digests": {"sha256": "fffff"},
+                            },
+                        ]
+                    },
+                }
+            )
+
+        raise NotImplementedError(url)
+
+    murlopen.side_effect = mocked_get
+
+    result = hashin.get_package_hashes(
+        package="hashin",
+        version="0.10",
+        algorithm="sha256",
+        index_url="https://pypi.internal.net/",
+    )
+
+    expected = {
+        "package": "hashin",
+        "version": "0.10",
+        "hashes": [{"hash": "ddddd"}, {"hash": "eeeee"}, {"hash": "fffff"}],
+    }
+
+    assert result == expected
+
+
 def test_get_package_hashes_package_not_found(murlopen):
     def mocked_get(url, **options):
         if url == "https://pypi.org/pypi/gobblygook/json":