Fixes to get pex to work on windows.

pex/common.py

@@ -28,18 +28,26 @@ def do_copy():
     shutil.copyfile(source, temp_dest)
     os.rename(temp_dest, dest)
 
-  try:
-    os.link(source, dest)
-  except OSError as e:
-    if e.errno == errno.EEXIST:
-      # File already exists.  If overwrite=True, write otherwise skip.
-      if overwrite:
+  # If the platform supports hard-linking, use that and fall back to copying.
+  # Windows does not support hard-linking.
+  if hasattr(os, 'link'):
+    try:
+      os.link(source, dest)
+    except OSError as e:
+      if e.errno == errno.EEXIST:
+        # File already exists.  If overwrite=True, write otherwise skip.
+        if overwrite:
+          do_copy()
+      elif e.errno == errno.EXDEV:
+        # Hard link across devices, fall back on copying
         do_copy()
-    elif e.errno == errno.EXDEV:
-      # Hard link across devices, fall back on copying
+      else:
+        raise
+  elif os.path.exists(dest):
+    if overwrite:
       do_copy()
-    else:
-      raise
+  else:
+    do_copy()
 
 
 # See http://stackoverflow.com/questions/2572172/referencing-other-modules-in-atexit
@@ -264,17 +272,8 @@ def link(self, src, dst, label=None):
     self._ensure_parent(dst)
     abs_src = src
     abs_dst = os.path.join(self.chroot, dst)
-    try:
-      os.link(abs_src, abs_dst)
-    except OSError as e:
-      if e.errno == errno.EEXIST:
-        # File already exists, skip XXX -- ensure target and dest are same?
-        pass
-      elif e.errno == errno.EXDEV:
-        # Hard link across devices, fall back on copying
-        shutil.copyfile(abs_src, abs_dst)
-      else:
-        raise
+    safe_copy(abs_src, abs_dst, overwrite=False)
+    # TODO: Ensure the target and dest are the same if the file already exists.
 
   def write(self, data, dst, label=None, mode='wb'):
     """Write data to ``chroot/dst`` with optional label.

pex/compiler.py

@@ -4,9 +4,9 @@
 from __future__ import absolute_import
 
 import subprocess
-import tempfile
 
 from .compatibility import to_bytes
+from .util import named_temporary_file
 
 
 _COMPILER_MAIN = """
@@ -78,7 +78,7 @@ def compile(self, root, relpaths):
     :returns: A list of relative paths of the compiled bytecode files.
     :raises: A :class:`Compiler.Error` if there was a problem bytecode compiling any of the files.
     """
-    with tempfile.NamedTemporaryFile() as fp:
+    with named_temporary_file() as fp:
       fp.write(to_bytes(_COMPILER_MAIN % {'root': root, 'relpaths': relpaths}, encoding='utf-8'))
       fp.flush()
       process = subprocess.Popen([self._interpreter.binary, fp.name],

pex/pex.py

@@ -96,14 +96,21 @@ def _extras_paths(cls):
       yield os.path.join(standard_lib, path)
 
   @classmethod
-  def _site_libs(cls):
+  def _get_site_packages(cls):
     try:
       from site import getsitepackages
-      site_libs = set(getsitepackages())
+      return set(getsitepackages())
     except ImportError:
-      site_libs = set()
+      return set()
+
+  @classmethod
+  def site_libs(cls):
+    site_libs = cls._get_site_packages()
     site_libs.update([sysconfig.get_python_lib(plat_specific=False),
                       sysconfig.get_python_lib(plat_specific=True)])
+    # On windows getsitepackages() returns the python stdlib too.
+    if sys.prefix in site_libs:
+      site_libs.remove(sys.prefix)
     real_site_libs = set(os.path.realpath(path) for path in site_libs)
     return site_libs | real_site_libs
 
@@ -190,7 +197,7 @@ def minimum_sys(cls):
     :returns: (sys.path, sys.path_importer_cache, sys.modules) tuple of a
       bare python installation.
     """
-    site_libs = set(cls._site_libs())
+    site_libs = set(cls.site_libs())
     for site_lib in site_libs:
       TRACER.log('Found site-library: %s' % site_lib)
     for extras_path in cls._extras_paths():

pex/util.py

@@ -7,6 +7,7 @@
 import errno
 import os
 import shutil
+import tempfile
 import uuid
 from hashlib import sha1
 from threading import Lock
@@ -200,3 +201,19 @@ def get(self, key, default=None):
   def store(self, key, value):
     with self._lock:
       self._data[key] = value
+
+
+@contextlib.contextmanager
+def named_temporary_file(*args, **kwargs):
+  """
+  Due to a bug in python (https://bugs.python.org/issue14243), we need
+  this to be able to use the temporary file without deleting it.
+  """
+  assert 'delete' not in kwargs
+  kwargs['delete'] = False
+  fp = tempfile.NamedTemporaryFile(*args, **kwargs)
+  try:
+    with fp:
+      yield fp
+  finally:
+    os.remove(fp.name)

tests/test_pex.py

@@ -2,17 +2,24 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 import os
+import sys
 import textwrap
 from types import ModuleType
 
 import pytest
+from twitter.common.contextutil import temporary_dir
 
-from pex.compatibility import to_bytes
+from pex.compatibility import nested, to_bytes
 from pex.installer import EggInstaller, WheelInstaller
 from pex.pex import PEX
 from pex.testing import make_installer, run_simple_pex_test
 from pex.util import DistributionHelper
 
+try:
+  from unittest import mock
+except ImportError:
+  import mock
+
 
 @pytest.mark.skipif('sys.version_info > (3,)')
 def test_pex_uncaught_exceptions():
@@ -120,6 +127,43 @@ def test_minimum_sys_modules():
   assert tainted_module.__path__ == ['good_path']
 
 
+def test_site_libs():
+  with nested(mock.patch.object(PEX, '_get_site_packages'), temporary_dir()) as (
+          mock_site_packages, tempdir):
+    site_packages = os.path.join(tempdir, 'site-packages')
+    os.mkdir(site_packages)
+    mock_site_packages.return_value = set([site_packages])
+    site_libs = PEX.site_libs()
+    assert site_packages in site_libs
+
+
+def test_site_libs_symlink():
+  with nested(mock.patch.object(PEX, '_get_site_packages'), temporary_dir()) as (
+          mock_site_packages, tempdir):
+    site_packages = os.path.join(tempdir, 'site-packages')
+    os.mkdir(site_packages)
+    site_packages_link = os.path.join(tempdir, 'site-packages-link')
+    os.symlink(site_packages, site_packages_link)
+    mock_site_packages.return_value = set([site_packages_link])
+
+    site_libs = PEX.site_libs()
+    assert os.path.realpath(site_packages) in site_libs
+    assert site_packages_link in site_libs
+
+
+def test_site_libs_excludes_prefix():
+  """Windows returns sys.prefix as part of getsitepackages(). Make sure to exclude it."""
+
+  with nested(mock.patch.object(PEX, '_get_site_packages'), temporary_dir()) as (
+          mock_site_packages, tempdir):
+    site_packages = os.path.join(tempdir, 'site-packages')
+    os.mkdir(site_packages)
+    mock_site_packages.return_value = set([site_packages, sys.prefix])
+    site_libs = PEX.site_libs()
+    assert site_packages in site_libs
+    assert sys.prefix not in site_libs
+
+
 @pytest.mark.parametrize('zip_safe', (False, True))
 @pytest.mark.parametrize('project_name', ('my_project', 'my-project'))
 @pytest.mark.parametrize('installer_impl', (EggInstaller, WheelInstaller))

tests/test_util.py

@@ -16,7 +16,7 @@
 from pex.installer import EggInstaller, WheelInstaller
 from pex.pex_builder import PEXBuilder
 from pex.testing import make_bdist, temporary_content, write_zipfile
-from pex.util import CacheHelper, DistributionHelper
+from pex.util import CacheHelper, DistributionHelper, named_temporary_file
 
 try:
   from unittest import mock
@@ -159,3 +159,16 @@ def test_access_zipped_assets_integration():
       pass
     assert output == 'accessed\n'
     assert po.returncode == 0
+
+
+def test_named_temporary_file():
+  name = ''
+  with named_temporary_file() as fp:
+    name = fp.name
+    fp.write(b'hi')
+    fp.flush()
+    assert os.path.exists(name)
+    with open(name) as new_fp:
+      assert new_fp.read() == 'hi'
+
+  assert not os.path.exists(name)