Improve failure modes for os.rename() as used in distribution caching.

pex/common.py

@@ -140,6 +140,21 @@ def safe_rmtree(directory):
     shutil.rmtree(directory, True)
 
 
+def rename_if_empty(src, dest, allowable_errors=(errno.EEXIST, errno.ENOTEMPTY)):
+  """Rename `src` to `dest` using `os.rename()`.
+
+  If an `OSError` with errno in `allowable_errors` is encountered during the rename, the `dest`
+  dir is left unchanged and the `src` directory will simply be removed.
+  """
+  try:
+    os.rename(src, dest)
+  except OSError as e:
+    if e.errno in allowable_errors:
+      safe_rmtree(src)
+    else:
+      raise
+
+
 def chmod_plus_x(path):
   """Equivalent of unix `chmod a+x path`"""
   path_mode = os.stat(path).st_mode

pex/util.py

@@ -4,7 +4,6 @@
 from __future__ import absolute_import
 
 import contextlib
-import errno
 import os
 import shutil
 import tempfile
@@ -14,7 +13,7 @@
 
 from pkg_resources import find_distributions, resource_isdir, resource_listdir, resource_string
 
-from .common import safe_mkdir, safe_mkdtemp, safe_open, safe_rmtree
+from .common import rename_if_empty, safe_mkdir, safe_mkdtemp, safe_open
 from .finders import register_finders
 
 
@@ -176,13 +175,8 @@ def cache_distribution(cls, zf, source, target_dir):
           with contextlib.closing(zf.open(name)) as zi:
             with safe_open(os.path.join(target_dir_tmp, target_name), 'wb') as fp:
               shutil.copyfileobj(zi, fp)
-      try:
-        os.rename(target_dir_tmp, target_dir)
-      except OSError as e:
-        if e.errno == errno.ENOTEMPTY:
-          safe_rmtree(target_dir_tmp)
-        else:
-          raise
+
+      rename_if_empty(target_dir_tmp, target_dir)
 
     dist = DistributionHelper.distribution_from_path(target_dir)
     assert dist is not None, 'Failed to cache distribution %s' % source

tests/test_common.py

@@ -0,0 +1,44 @@
+# Copyright 2016 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+import errno
+import os
+from contextlib import contextmanager
+
+import pytest
+
+from pex.common import rename_if_empty
+
+try:
+  from unittest import mock
+except ImportError:
+  import mock
+
+
+@contextmanager
+def maybe_raises(exception=None):
+  @contextmanager
+  def noop():
+    yield
+
+  with (noop() if exception is None else pytest.raises(exception)):
+    yield
+
+
+def rename_if_empty_test(errno, expect_raises=None):
+  with mock.patch('os.rename', spec_set=True, autospec=True) as mock_rename:
+    mock_rename.side_effect = OSError(errno, os.strerror(errno))
+    with maybe_raises(expect_raises):
+      rename_if_empty('from.dir', 'to.dir')
+
+
+def test_rename_if_empty_eexist():
+  rename_if_empty_test(errno.EEXIST)
+
+
+def test_rename_if_empty_enotempty():
+  rename_if_empty_test(errno.ENOTEMPTY)
+
+
+def test_rename_if_empty_eperm():
+  rename_if_empty_test(errno.EPERM, expect_raises=OSError)