bpo-35854: Fix EnvBuilder and --symlinks in venv on Windows (GH-11700)

2 files changed, 44 insertions, 18 deletions

diff --git a/Lib/test/test_venv.py b/Lib/test/test_venv.py
index 34c2234..6096b9d 100644
--- a/Lib/test/test_venv.py
+++ b/Lib/test/test_venv.py
@@ -243,7 +243,6 @@ class BasicTest(BaseTest):
             self.assertIn('include-system-site-packages = %s\n' % s, data)
 
     @unittest.skipUnless(can_symlink(), 'Needs symlinks')
-    @unittest.skipIf(os.name == 'nt', 'Symlinks are never used on Windows')
     def test_symlinking(self):
         """
         Test symlinking works as expected
diff --git a/Lib/venv/__init__.py b/Lib/venv/__init__.py
index 5438b0d..8f9e313 100644
--- a/Lib/venv/__init__.py
+++ b/Lib/venv/__init__.py
@@ -64,11 +64,10 @@ class EnvBuilder:
         self.system_site_packages = False
         self.create_configuration(context)
         self.setup_python(context)
-        if not self.upgrade:
-            self.setup_scripts(context)
         if self.with_pip:
             self._setup_pip(context)
         if not self.upgrade:
+            self.setup_scripts(context)
             self.post_setup(context)
         if true_system_site_packages:
             # We had set it to False before, now
@@ -176,6 +175,23 @@ class EnvBuilder:
                 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))
+                if basename.endswith('_d'):
+                    ext = '_d' + ext
+                    basename = basename[:-2]
+                if sysconfig.is_python_build(True):
+                    if basename == 'python':
+                        basename = 'venvlauncher'
+                    elif basename == 'pythonw':
+                        basename = 'venvwlauncher'
+                    scripts = os.path.dirname(src)
+                else:
+                    scripts = os.path.join(os.path.dirname(__file__), "scripts", "nt")
+                src = os.path.join(scripts, basename + ext)
+
             shutil.copyfile(src, dst)
 
     def setup_python(self, context):
@@ -202,23 +218,31 @@ class EnvBuilder:
                     if not os.path.islink(path):
                         os.chmod(path, 0o755)
         else:
-            # For normal cases, the venvlauncher will be copied from
-            # our scripts folder. For builds, we need to copy it
-            # manually.
-            if sysconfig.is_python_build(True):
-                suffix = '.exe'
-                if context.python_exe.lower().endswith('_d.exe'):
-                    suffix = '_d.exe'
-
-                src = os.path.join(dirname, "venvlauncher" + suffix)
-                dst = os.path.join(binpath, context.python_exe)
-                copier(src, dst)
+            if self.symlinks:
+                # For symlinking, we need a complete copy of the root directory
+                # If symlinks fail, you'll get unnecessary copies of files, but
+                # we assume that if you've opted into symlinks on Windows then
+                # you know what you're doing.
+                suffixes = [
+                    f for f in os.listdir(dirname) if
+                    os.path.normcase(os.path.splitext(f)[1]) in ('.exe', '.dll')
+                ]
+                if sysconfig.is_python_build(True):
+                    suffixes = [
+                        f for f in suffixes if
+                        os.path.normcase(f).startswith(('python', 'vcruntime'))
+                    ]
+            else:
+                suffixes = ['python.exe', 'python_d.exe', 'pythonw.exe',
+                            'pythonw_d.exe']
 
-                src = os.path.join(dirname, "venvwlauncher" + suffix)
-                dst = os.path.join(binpath, "pythonw" + suffix)
-                copier(src, dst)
+            for suffix in suffixes:
+                src = os.path.join(dirname, suffix)
+                if os.path.exists(src):
+                    copier(src, os.path.join(binpath, suffix))
 
-                # copy init.tcl over
+            if sysconfig.is_python_build(True):
+                # copy init.tcl
                 for root, dirs, files in os.walk(context.python_dir):
                     if 'init.tcl' in files:
                         tcldir = os.path.basename(root)
@@ -304,6 +328,9 @@ class EnvBuilder:
                         dirs.remove(d)
                 continue # ignore files in top level
             for f in files:
+                if (os.name == 'nt' and f.startswith('python')
+                        and f.endswith(('.exe', '.pdb'))):
+                    continue
                 srcfile = os.path.join(root, f)
                 suffix = root[plen:].split(os.sep)[2:]
                 if not suffix: