Add test.support.PythonSymlink to encapsulate platform-specific logic…

… for linking sys.executable
python · zooba · Jun 29, 2019 · Jun 26, 2019 · Jun 27, 2019 · Jun 27, 2019
commit 07f17503a676978f3a7f1394c0f522c808c6539c
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
@@ -12,6 +12,7 @@
 import fnmatch
 import functools
 import gc
+import glob
 import importlib
 import importlib.util
 import io
@@ -2500,6 +2501,84 @@ def skip_unless_symlink(test):
     msg = "Requires functional symlink implementation"
     return test if ok else unittest.skip(msg)(test)
 
+class PythonSymlink:
+    """Creates a symlink for the current Python executable"""
+    def __init__(self, link=None):
+        self.link = link or os.path.abspath(TESTFN)
+        self._linked = []
+        self.real = os.path.realpath(sys.executable)
+        self._also_link = []
+
+        self._env = None
+
+        self._platform_specific()
+
+    def _platform_specific(self):
+        pass
+
+    if sys.platform == "win32":
+        def _platform_specific(self):
+            import _winapi
+
+            if os.path.lexists(self.real) and not os.path.exists(self.real):
+                # App symlink appears to not exist, but we want the
+                # real executable here anyway
+                self.real = _winapi.GetModuleFileName(0)
+
+            dll = _winapi.GetModuleFileName(sys.dllhandle)
+            src_dir = os.path.dirname(dll)
+            dest_dir = os.path.dirname(self.link)
+            self._also_link.append((
+                dll,
+                os.path.join(dest_dir, os.path.basename(dll))
+            ))
+            for runtime in glob.glob(os.path.join(src_dir, "vcruntime*.dll")):
+                self._also_link.append((
+                    runtime,
+                    os.path.join(dest_dir, os.path.basename(runtime))
+                ))
+
+            self._env = {k.upper(): os.getenv(k) for k in os.environ}
+            self._env["PYTHONHOME"] = os.path.dirname(self.real)
+            if sysconfig.is_python_build(True):
+                self._env["PYTHONPATH"] = os.path.dirname(os.__file__)
+
+    def __enter__(self):
+        os.symlink(self.real, self.link)
+        self._linked.append(self.link)
+        for real, link in self._also_link:
+            os.symlink(real, link)
+            self._linked.append(link)
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        for link in self._linked:
+            try:
+                os.remove(link)
+            except IOError as ex:
+                if verbose:
+                    print("failed to clean up {}: {}".format(link, ex))
+
+    def _call(self, python, args, env, returncode):
+        cmd = [python, *args]
+        p = subprocess.Popen(cmd, stdout=subprocess.PIPE,
+                             stderr=subprocess.PIPE, env=env)
+        r = p.communicate()
+        if p.returncode != returncode:
+            if verbose:
+                print(repr(r[0]))
+                print(repr(r[1]), file=sys.stderr)
+            raise RuntimeError(
+                'unexpected return code: {0} (0x{0:08X})'.format(p.returncode))
+        return r
+
+    def call_real(self, *args, returncode=0):
+        return self._call(self.real, args, None, returncode)
+
+    def call_link(self, *args, returncode=0):
+        return self._call(self.link, args, self._env, returncode)
+
+
 _can_xattr = None
 def can_xattr():
     global _can_xattr

diff --git a/Lib/test/test_httpservers.py b/Lib/test/test_httpservers.py
@@ -610,9 +610,10 @@ def setUp(self):
 
         # The shebang line should be pure ASCII: use symlink if possible.
         # See issue #7668.
+        self._pythonexe_symlink = None
         if support.can_symlink():
             self.pythonexe = os.path.join(self.parent_dir, 'python')
-            os.symlink(sys.executable, self.pythonexe)
+            self._pythonexe_symlink = support.PythonSymlink(self.pythonexe).__enter__()
         else:
             self.pythonexe = sys.executable
 
@@ -655,8 +656,8 @@ def setUp(self):
     def tearDown(self):
         try:
             os.chdir(self.cwd)
-            if self.pythonexe != sys.executable:
-                os.remove(self.pythonexe)
+            if self._pythonexe_symlink:
+                self._pythonexe_symlink.__exit__(None, None, None)
             if self.nocgi_path:
                 os.remove(self.nocgi_path)
             if self.file1_path:

diff --git a/Lib/test/test_platform.py b/Lib/test/test_platform.py
@@ -20,44 +20,9 @@ def test_architecture(self):
 
     @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(real), env.get("PATH", ""))
-            env["PYTHONHOME"] = os.path.dirname(real)
-            if sysconfig.is_python_build(True):
-                env["PYTHONPATH"] = os.path.dirname(os.__file__)
-
-        def get(python, env=None):
-            cmd = [python, '-c',
-                'import platform; print(platform.architecture())']
-            p = subprocess.Popen(cmd, stdout=subprocess.PIPE,
-                                 stderr=subprocess.PIPE, env=env)
-            r = p.communicate()
-            if p.returncode:
-                print(repr(r[0]))
-                print(repr(r[1]), file=sys.stderr)
-                self.fail('unexpected return code: {0} (0x{0:08X})'
-                          .format(p.returncode))
-            return r
-
-        try:
-            self.assertEqual(get(sys.executable), get(link, env=env))
-        finally:
-            os.remove(link)
+        with support.PythonSymlink() as py:
+            cmd = "-c", "import platform; print(platform.architecture())"
+            self.assertEqual(py.call_real(*cmd), py.call_link(*cmd))
 
     def test_platform(self):
         for aliased in (False, True):

diff --git a/Lib/test/test_sysconfig.py b/Lib/test/test_sysconfig.py
@@ -6,7 +6,8 @@
 from copy import copy
 
 from test.support import (import_module, TESTFN, unlink, check_warnings,
-                          captured_stdout, skip_unless_symlink, change_cwd)
+                          captured_stdout, skip_unless_symlink, change_cwd,
+                          PythonSymlink)
 
 import sysconfig
 from sysconfig import (get_paths, get_platform, get_config_vars,
@@ -232,47 +233,10 @@ def test_get_scheme_names(self):
         self.assertEqual(get_scheme_names(), wanted)
 
     @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(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(real)
-            if sysconfig.is_python_build(True):
-                env["PYTHONPATH"] = os.path.dirname(os.__file__)
-
-        # Issue 7880
-        def get(python, env=None):
-            cmd = [python, '-c',
-                   'import sysconfig; print(sysconfig.get_platform())']
-            p = subprocess.Popen(cmd, stdout=subprocess.PIPE,
-                                 stderr=subprocess.PIPE, env=env)
-            out, err = p.communicate()
-            if p.returncode:
-                print((out, err))
-                self.fail('Non-zero return code {0} (0x{0:08X})'
-                            .format(p.returncode))
-            return out, err
-
-        try:
-            self.assertEqual(get(real), get(link, env))
-        finally:
-            unlink(link)
+    def test_symlink(self): # Issue 7880
+        with PythonSymlink() as py:
+            cmd = "-c", "import sysconfig; print(sysconfig.get_platform())"
+            self.assertEqual(py.call_real(*cmd), py.call_link(*cmd))
 
     def test_user_similar(self):
         # Issue #8759: make sure the posix scheme for the users