Fix a bug where pip-compile would ignore some of versions if prefer binary distributions

piptools/repositories/pypi.py

@@ -123,15 +123,16 @@ def find_best_match(self, ireq, prereleases=None):
             return ireq  # return itself as the best match
 
         all_candidates = self.find_all_candidates(ireq.name)
-        candidates_by_version = lookup_table(
-            all_candidates, key=lambda c: c.version, unique=True
-        )
+        candidates_by_version = lookup_table(all_candidates, key=lambda c: c.version)
         matching_versions = ireq.specifier.filter(
             (candidate.version for candidate in all_candidates), prereleases=prereleases
         )
 
-        # Reuses pip's internal candidate sort key to sort
-        matching_candidates = [candidates_by_version[ver] for ver in matching_versions]
+        matching_candidates = list(
+            itertools.chain.from_iterable(
+                candidates_by_version[ver] for ver in matching_versions
+            )
+        )
         if not matching_candidates:
             raise NoCandidateFound(ireq, all_candidates, self.finder)
 

tests/conftest.py

@@ -1,7 +1,9 @@
 import json
 import os
+import sys
 from contextlib import contextmanager
 from functools import partial
+from subprocess import check_call
 from textwrap import dedent
 
 import pytest
@@ -221,3 +223,82 @@ def pip_with_index_conf(make_pip_conf):
             )
         )
     )
+
+
+@pytest.fixture
+def make_package(tmp_path):
+    """
+    Make a package from a given name, version and list of required packages.
+    """
+
+    def _make_package(name, version="0.1", install_requires=None):
+        if install_requires is None:
+            install_requires = []
+
+        install_requires_str = "[{}]".format(
+            ",".join("{!r}".format(package) for package in install_requires)
+        )
+
+        package_dir = tmp_path / "packages" / name / version
+        package_dir.mkdir(parents=True)
+
+        setup_file = str(package_dir / "setup.py")
+        with open(setup_file, "w") as fp:
+            fp.write(
+                dedent(
+                    """\
+                    from setuptools import setup
+                    setup(
+                        name={name!r},
+                        version={version!r},
+                        install_requires={install_requires_str},
+                    )
+                    """.format(
+                        name=name,
+                        version=version,
+                        install_requires_str=install_requires_str,
+                    )
+                )
+            )
+        return package_dir
+
+    return _make_package
+
+
+@pytest.fixture
+def run_setup_file():
+    """
+    Run a setup.py file from a given package dir.
+    """
+
+    def _make_wheel(package_dir_path, *args):
+        setup_file = str(package_dir_path / "setup.py")
+        return check_call((sys.executable, setup_file) + args)  # nosec
+
+    return _make_wheel
+
+
+@pytest.fixture
+def make_wheel(run_setup_file):
+    """
+    Make a wheel distribution from a given package dir.
+    """
+
+    def _make_wheel(package_dir, dist_dir, *args):
+        return run_setup_file(
+            package_dir, "bdist_wheel", "--dist-dir", str(dist_dir), *args
+        )
+
+    return _make_wheel
+
+
+@pytest.fixture
+def make_sdist(run_setup_file):
+    """
+    Make a source distribution from a given package dir.
+    """
+
+    def _make_sdist(package_dir, dist_dir, *args):
+        return run_setup_file(package_dir, "sdist", "--dist-dir", str(dist_dir), *args)
+
+    return _make_sdist

tests/test_cli_compile.py

@@ -1146,3 +1146,75 @@ def test_preserve_compiled_prerelease_version(pip_conf, runner):
 
     assert out.exit_code == 0, out
     assert "small-fake-a==0.3b1" in out.stderr.splitlines()
+
+
+def test_prefer_binary_dist(
+    pip_conf, make_package, make_sdist, make_wheel, tmpdir, runner
+):
+    """
+    Test pip-compile chooses a correct version of a package with
+    a binary distribution when PIP_PREFER_BINARY environment variable is on.
+    """
+    dists_dir = tmpdir / "dists"
+
+    # Make first-package==1.0 and wheels
+    first_package_v1 = make_package(name="first-package", version="1.0")
+    make_wheel(first_package_v1, dists_dir)
+
+    # Make first-package==2.0 and sdists
+    first_package_v2 = make_package(name="first-package", version="2.0")
+    make_sdist(first_package_v2, dists_dir)
+
+    # Make second-package==1.0 which depends on first-package, and wheels
+    second_package_v1 = make_package(
+        name="second-package", version="1.0", install_requires=["first-package"]
+    )
+    make_wheel(second_package_v1, dists_dir)
+
+    with open("requirements.in", "w") as req_in:
+        req_in.write("second-package")
+
+    out = runner.invoke(
+        cli,
+        ["--no-annotate", "--find-links", str(dists_dir)],
+        env={"PIP_PREFER_BINARY": "1"},
+    )
+
+    assert out.exit_code == 0, out
+    assert "first-package==1.0" in out.stderr.splitlines(), out.stderr
+    assert "second-package==1.0" in out.stderr.splitlines(), out.stderr
+
+
+@pytest.mark.parametrize("prefer_binary", (True, False))
+def test_prefer_binary_dist_even_there_is_source_dists(
+    pip_conf, make_package, make_sdist, make_wheel, tmpdir, runner, prefer_binary
+):
+    """
+    Test pip-compile chooses a correct version of a package with a binary distribution
+    (despite a source dist existing) when PIP_PREFER_BINARY environment variable is on
+    or off.
+
+    Regression test for issue GH-1118.
+    """
+    dists_dir = tmpdir / "dists"
+
+    # Make first version of package with only wheels
+    package_v1 = make_package(name="test-package", version="1.0")
+    make_wheel(package_v1, dists_dir)
+
+    # Make seconds version with wheels and sdists
+    package_v2 = make_package(name="test-package", version="2.0")
+    make_wheel(package_v2, dists_dir)
+    make_sdist(package_v2, dists_dir)
+
+    with open("requirements.in", "w") as req_in:
+        req_in.write("test-package")
+
+    out = runner.invoke(
+        cli,
+        ["--no-annotate", "--find-links", str(dists_dir)],
+        env={"PIP_PREFER_BINARY": str(int(prefer_binary))},
+    )
+
+    assert out.exit_code == 0, out
+    assert "test-package==2.0" in out.stderr.splitlines(), out.stderr