From 4d11b025f188d65d007ad89f54996e6ffb320958 Mon Sep 17 00:00:00 2001
From: Todd Gardner <todd.gardner@gmail.com>
Date: Thu, 6 Apr 2017 14:38:17 -0400
Subject: [PATCH] Allow the caching and static iterator to respect
 allow_prereleases (#373)

---
 pex/resolver.py        | 21 ++++++++++-----
 tests/test_resolver.py | 61 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 7 deletions(-)

diff --git a/pex/resolver.py b/pex/resolver.py
index 41f0e0777..1e842837a 100644
--- a/pex/resolver.py
+++ b/pex/resolver.py
@@ -35,12 +35,13 @@ class Unsatisfiable(Exception):
 class StaticIterator(IteratorInterface):
   """An iterator that iterates over a static list of packages."""
 
-  def __init__(self, packages):
+  def __init__(self, packages, allow_prereleases=None):
     self._packages = packages
+    self._allow_prereleases = allow_prereleases
 
   def iter(self, req):
     for package in self._packages:
-      if package.satisfies(req):
+      if package.satisfies(req, allow_prereleases=self._allow_prereleases):
         yield package
 
 
@@ -150,13 +151,15 @@ def filter_packages_by_interpreter(cls, packages, interpreter, platform):
     return [package for package in packages
         if package.compatible(interpreter.identity, platform)]
 
-  def __init__(self, interpreter=None, platform=None):
+  def __init__(self, allow_prereleases=None, interpreter=None, platform=None):
     self._interpreter = interpreter or PythonInterpreter.get()
     self._platform = platform or Platform.current()
+    self._allow_prereleases = allow_prereleases
 
   def package_iterator(self, resolvable, existing=None):
     if existing:
-      existing = resolvable.compatible(StaticIterator(existing))
+      existing = resolvable.compatible(
+        StaticIterator(existing, allow_prereleases=self._allow_prereleases))
     else:
       existing = resolvable.packages()
     return self.filter_packages_by_interpreter(existing, self._interpreter, self._platform)
@@ -239,7 +242,8 @@ def __init__(self, cache, cache_ttl, *args, **kw):
 
   # Short-circuiting package iterator.
   def package_iterator(self, resolvable, existing=None):
-    iterator = Iterator(fetchers=[Fetcher([self.__cache])])
+    iterator = Iterator(fetchers=[Fetcher([self.__cache])],
+                        allow_prereleases=self._allow_prereleases)
     packages = self.filter_packages_by_interpreter(
       resolvable.compatible(iterator),
       self._interpreter,
@@ -350,8 +354,11 @@ def resolve(
   )
 
   if cache:
-    resolver = CachingResolver(cache, cache_ttl, interpreter=interpreter, platform=platform)
+    resolver = CachingResolver(
+      cache, cache_ttl,
+      allow_prereleases=allow_prereleases, interpreter=interpreter, platform=platform)
   else:
-    resolver = Resolver(interpreter=interpreter, platform=platform)
+    resolver = Resolver(
+      allow_prereleases=allow_prereleases, interpreter=interpreter, platform=platform)
 
   return resolver.resolve(resolvables_from_iterable(requirements, builder))
diff --git a/tests/test_resolver.py b/tests/test_resolver.py
index e08414b60..26498ec8e 100644
--- a/tests/test_resolver.py
+++ b/tests/test_resolver.py
@@ -102,6 +102,67 @@ def assert_resolve(expected_version, **resolve_kwargs):
     assert_resolve('3.0.0rc3', allow_prereleases=True)
 
 
+def test_resolve_prereleases_cached():
+  stable_dep = make_sdist(name='dep', version='2.0.0')
+  prerelease_dep = make_sdist(name='dep', version='3.0.0rc3')
+
+  with temporary_dir() as td:
+    for sdist in (stable_dep, prerelease_dep):
+      safe_copy(sdist, os.path.join(td, os.path.basename(sdist)))
+    fetchers = [Fetcher([td])]
+
+    with temporary_dir() as cd:
+      def assert_resolve(dep, expected_version, **resolve_kwargs):
+        dists = resolve(
+          [dep], cache=cd, cache_ttl=1000, **resolve_kwargs)
+        assert 1 == len(dists)
+        dist = dists[0]
+        assert expected_version == dist.version
+
+      Crawler.reset_cache()
+
+      # First do a run to load it into the cache.
+      assert_resolve('dep>=1,<4', '3.0.0rc3', allow_prereleases=True, fetchers=fetchers)
+
+      # This simulates running from another pex command. The Crawler cache actually caches an empty
+      # cache so this fails in the same "process".
+      Crawler.reset_cache()
+
+      # Now assert that we can get it from the cache by removing the source.
+      assert_resolve('dep>=1,<4', '3.0.0rc3', allow_prereleases=True, fetchers=[])
+
+      # It should also be able to resolve without allow_prereleases, if explicitly requested.
+      Crawler.reset_cache()
+      assert_resolve('dep>=1.rc1,<4', '3.0.0rc3', fetchers=[])
+
+
+def test_resolve_prereleases_multiple_set():
+  stable_dep = make_sdist(name='dep', version='2.0.0')
+  prerelease_dep1 = make_sdist(name='dep', version='3.0.0rc3')
+  prerelease_dep2 = make_sdist(name='dep', version='3.0.0rc4')
+  prerelease_dep3 = make_sdist(name='dep', version='3.0.0rc5')
+
+  with temporary_dir() as td:
+    for sdist in (stable_dep, prerelease_dep1, prerelease_dep2, prerelease_dep3):
+      safe_copy(sdist, os.path.join(td, os.path.basename(sdist)))
+    fetchers = [Fetcher([td])]
+
+    def assert_resolve(expected_version, **resolve_kwargs):
+      dists = resolve(
+        [
+          'dep>=3.0.0rc1',
+          'dep==3.0.0rc4',
+        ],
+        fetchers=fetchers, **resolve_kwargs)
+      assert 1 == len(dists)
+      dist = dists[0]
+      assert expected_version == dist.version
+
+    # This should resolve with explicit prerelease being set or implicitly.
+    assert_resolve('3.0.0rc4', allow_prereleases=True)
+    assert_resolve('3.0.0rc4')
+
+
 def test_resolvable_set():
   builder = ResolverOptionsBuilder()
   rs = _ResolvableSet()