bpo-37369: Fix venv and test symlinking (GH-14456)

4 files changed, 96 insertions, 47 deletions

diff --git a/Lib/test/test_platform.py b/Lib/test/test_platform.py
index 010ed6c634..d91e978a79 100644
--- a/Lib/test/test_platform.py
+++ b/Lib/test/test_platform.py
@@ -16,14 +16,24 @@ class PlatformTest(unittest.TestCase):
 
     @support.skip_unless_symlink
     def test_architecture_via_symlink(self): # issue3762
+        if sys.platform == "win32" and not os.path.exists(sys.executable):
+            # App symlink appears to not exist, but we want the
+            # real executable here anyway
+            import _winapi
+            real = _winapi.GetModuleFileName(0)
+        else:
+            real = os.path.realpath(sys.executable)
+        link = os.path.abspath(support.TESTFN)
+        os.symlink(real, link)
+
         # On Windows, the EXE needs to know where pythonXY.dll and *.pyd is at
         # so we add the directory to the path, PYTHONHOME and PYTHONPATH.
         env = None
         if sys.platform == "win32":
             env = {k.upper(): os.environ[k] for k in os.environ}
             env["PATH"] = "{};{}".format(
-                os.path.dirname(sys.executable), env.get("PATH", ""))
-            env["PYTHONHOME"] = os.path.dirname(sys.executable)
+                os.path.dirname(real), env.get("PATH", ""))
+            env["PYTHONHOME"] = os.path.dirname(real)
             if sysconfig.is_python_build(True):
                 env["PYTHONPATH"] = os.path.dirname(os.__file__)
 
@@ -40,11 +50,8 @@ class PlatformTest(unittest.TestCase):
                           .format(p.returncode))
             return r
 
-        real = os.path.realpath(sys.executable)
-        link = os.path.abspath(support.TESTFN)
-        os.symlink(real, link)
         try:
-            self.assertEqual(get(real), get(link, env=env))
+            self.assertEqual(get(sys.executable), get(link, env=env))
         finally:
             os.remove(link)
 
@@ -280,6 +287,11 @@ class PlatformTest(unittest.TestCase):
            os.path.exists(sys.executable+'.exe'):
             # Cygwin horror
             executable = sys.executable + '.exe'
+        elif sys.platform == "win32" and not os.path.exists(sys.executable):
+            # App symlink appears to not exist, but we want the
+            # real executable here anyway
+            import _winapi
+            executable = _winapi.GetModuleFileName(0)
         else:
             executable = sys.executable
         res = platform.libc_ver(executable)
diff --git a/Lib/test/test_sysconfig.py b/Lib/test/test_sysconfig.py
index 1b1929885e..51bef19000 100644
--- a/Lib/test/test_sysconfig.py
+++ b/Lib/test/test_sysconfig.py
@@ -233,16 +233,26 @@ class TestSysConfig(unittest.TestCase):
 
     @skip_unless_symlink
     def test_symlink(self):
+        if sys.platform == "win32" and not os.path.exists(sys.executable):
+            # App symlink appears to not exist, but we want the
+            # real executable here anyway
+            import _winapi
+            real = _winapi.GetModuleFileName(0)
+        else:
+            real = os.path.realpath(sys.executable)
+        link = os.path.abspath(TESTFN)
+        os.symlink(real, link)
+
         # On Windows, the EXE needs to know where pythonXY.dll is at so we have
         # to add the directory to the path.
         env = None
         if sys.platform == "win32":
             env = {k.upper(): os.environ[k] for k in os.environ}
             env["PATH"] = "{};{}".format(
-                os.path.dirname(sys.executable), env.get("PATH", ""))
+                os.path.dirname(real), env.get("PATH", ""))
             # Requires PYTHONHOME as well since we locate stdlib from the
             # EXE path and not the DLL path (which should be fixed)
-            env["PYTHONHOME"] = os.path.dirname(sys.executable)
+            env["PYTHONHOME"] = os.path.dirname(real)
             if sysconfig.is_python_build(True):
                 env["PYTHONPATH"] = os.path.dirname(os.__file__)
 
@@ -258,9 +268,6 @@ class TestSysConfig(unittest.TestCase):
                 self.fail('Non-zero return code {0} (0x{0:08X})'
                             .format(p.returncode))
             return out, err
-        real = os.path.realpath(sys.executable)
-        link = os.path.abspath(TESTFN)
-        os.symlink(real, link)
         try:
             self.assertEqual(get(real), get(link, env))
         finally:
diff --git a/Lib/test/test_venv.py b/Lib/test/test_venv.py
index c3ccb92913..67f9f46e65 100644
--- a/Lib/test/test_venv.py
+++ b/Lib/test/test_venv.py
@@ -58,6 +58,12 @@ class BaseTest(unittest.TestCase):
             self.include = 'include'
         executable = getattr(sys, '_base_executable', sys.executable)
         self.exe = os.path.split(executable)[-1]
+        if (sys.platform == 'win32'
+            and os.path.lexists(executable)
+            and not os.path.exists(executable)):
+            self.cannot_link_exe = True
+        else:
+            self.cannot_link_exe = False
 
     def tearDown(self):
         rmtree(self.env_dir)
@@ -248,7 +254,12 @@ class BasicTest(BaseTest):
             # symlinked to 'python3.3' in the env, even when symlinking in
             # general isn't wanted.
             if usl:
-                self.assertTrue(os.path.islink(fn))
+                if self.cannot_link_exe:
+                    # Symlinking is skipped when our executable is already a
+                    # special app symlink
+                    self.assertFalse(os.path.islink(fn))
+                else:
+                    self.assertTrue(os.path.islink(fn))
 
     # If a venv is created from a source build and that venv is used to
     # run the test, the pyvenv.cfg in the venv created in the test will
diff --git a/Lib/venv/__init__.py b/Lib/venv/__init__.py
index 95c05486af..c4540827a9 100644
--- a/Lib/venv/__init__.py
+++ b/Lib/venv/__init__.py
@@ -155,47 +155,66 @@ class EnvBuilder:
             f.write('include-system-site-packages = %s\n' % incl)
             f.write('version = %d.%d.%d\n' % sys.version_info[:3])
 
-    def symlink_or_copy(self, src, dst, relative_symlinks_ok=False):
-        """
-        Try symlinking a file, and if that fails, fall back to copying.
-        """
-        force_copy = not self.symlinks
-        if not force_copy:
-            try:
-                if not os.path.islink(dst): # can't link to itself!
+    if os.name != 'nt':
+        def symlink_or_copy(self, src, dst, relative_symlinks_ok=False):
+            """
+            Try symlinking a file, and if that fails, fall back to copying.
+            """
+            force_copy = not self.symlinks
+            if not force_copy:
+                try:
+                    if not os.path.islink(dst): # can't link to itself!
+                        if relative_symlinks_ok:
+                            assert os.path.dirname(src) == os.path.dirname(dst)
+                            os.symlink(os.path.basename(src), dst)
+                        else:
+                            os.symlink(src, dst)
+                except Exception:   # may need to use a more specific exception
+                    logger.warning('Unable to symlink %r to %r', src, dst)
+                    force_copy = True
+            if force_copy:
+                shutil.copyfile(src, dst)
+    else:
+        def symlink_or_copy(self, src, dst, relative_symlinks_ok=False):
+            """
+            Try symlinking a file, and if that fails, fall back to copying.
+            """
+            bad_src = os.path.lexists(src) and not os.path.exists(src)
+            if self.symlinks and not bad_src and not os.path.islink(dst):
+                try:
                     if relative_symlinks_ok:
                         assert os.path.dirname(src) == os.path.dirname(dst)
                         os.symlink(os.path.basename(src), dst)
                     else:
                         os.symlink(src, dst)
-            except Exception:   # may need to use a more specific exception
-                logger.warning('Unable to symlink %r to %r', src, dst)
-                force_copy = True
-        if force_copy:
-            if os.name == 'nt':
-                # On Windows, we rewrite symlinks to our base python.exe into
-                # copies of venvlauncher.exe
-                basename, ext = os.path.splitext(os.path.basename(src))
-                srcfn = os.path.join(os.path.dirname(__file__),
-                                     "scripts",
-                                     "nt",
-                                     basename + ext)
-                # Builds or venv's from builds need to remap source file
-                # locations, as we do not put them into Lib/venv/scripts
-                if sysconfig.is_python_build(True) or not os.path.isfile(srcfn):
-                    if basename.endswith('_d'):
-                        ext = '_d' + ext
-                        basename = basename[:-2]
-                    if basename == 'python':
-                        basename = 'venvlauncher'
-                    elif basename == 'pythonw':
-                        basename = 'venvwlauncher'
-                    src = os.path.join(os.path.dirname(src), basename + ext)
-                else:
-                    src = srcfn
-                if not os.path.exists(src):
-                    logger.warning('Unable to copy %r', src)
                     return
+                except Exception:   # may need to use a more specific exception
+                    logger.warning('Unable to symlink %r to %r', src, dst)
+
+            # On Windows, we rewrite symlinks to our base python.exe into
+            # copies of venvlauncher.exe
+            basename, ext = os.path.splitext(os.path.basename(src))
+            srcfn = os.path.join(os.path.dirname(__file__),
+                                 "scripts",
+                                 "nt",
+                                 basename + ext)
+            # Builds or venv's from builds need to remap source file
+            # locations, as we do not put them into Lib/venv/scripts
+            if sysconfig.is_python_build(True) or not os.path.isfile(srcfn):
+                if basename.endswith('_d'):
+                    ext = '_d' + ext
+                    basename = basename[:-2]
+                if basename == 'python':
+                    basename = 'venvlauncher'
+                elif basename == 'pythonw':
+                    basename = 'venvwlauncher'
+                src = os.path.join(os.path.dirname(src), basename + ext)
+            else:
+                src = srcfn
+            if not os.path.exists(src):
+                if not bad_src:
+                    logger.warning('Unable to copy %r', src)
+                return
 
             shutil.copyfile(src, dst)