From e88b5f16416d645c80c7e07f6ddf091e802ef130 Mon Sep 17 00:00:00 2001
From: Mike Kaplinskiy <mkaplinskiy@twitter.com>
Date: Mon, 11 Jan 2016 16:50:06 -0800
Subject: [PATCH] Fixes to get pex to work on windows.

The fixes are:
 - os.link doesn't exist on windows. Always copy instead.
 - NamedTemporaryFile doesn't work correctly on windows (see https://bugs.python.org/issue14243)
 - sys.prefix is part of site.getsitepackages() on windows. Don't remove it.
---
 pex/common.py   | 41 ++++++++++++++++++++---------------------
 pex/compiler.py |  4 ++--
 pex/pex.py      |  3 +++
 pex/util.py     | 17 +++++++++++++++++
 4 files changed, 42 insertions(+), 23 deletions(-)

diff --git a/pex/common.py b/pex/common.py
index 89f19419d..acdad9135 100644
--- a/pex/common.py
+++ b/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.
diff --git a/pex/compiler.py b/pex/compiler.py
index 6346d0ff2..d301de8ff 100644
--- a/pex/compiler.py
+++ b/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],
diff --git a/pex/pex.py b/pex/pex.py
index 111ef0f8a..7afd617cd 100644
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -104,6 +104,9 @@ def _site_libs(cls):
       site_libs = set()
     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
 
diff --git a/pex/util.py b/pex/util.py
index d9dbe2ef2..d5e3c494b 100644
--- a/pex/util.py
+++ b/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)