make an exe stub for windows

asottile · asottile · commit fbd529204bac · 2020-09-28T18:37:10.000-07:00
diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -1,5 +1,6 @@
 covdefaults
 coverage
+distlib
 pytest
 pytest-env
 re-assert
diff --git a/testing/zipapp/entry b/testing/zipapp/entry
@@ -33,7 +33,8 @@ def _ensure_cache(zipf: zipfile.ZipFile, cache_key: str) -> str:
         try:
             zipf.extractall(tmpdir)
             # zip doesn't maintain permissions
-            _make_executable(os.path.join(tmpdir, 'fakepython'))
+            _make_executable(os.path.join(tmpdir, 'python'))
+            _make_executable(os.path.join(tmpdir, 'python.exe'))
             os.rename(tmpdir, cache_dest)
         except BaseException:
             shutil.rmtree(tmpdir)
@@ -49,8 +50,12 @@ def main() -> int:
 
         cache_dest = _ensure_cache(zipf, cache_key)
 
-    fakepython = os.path.join(cache_dest, 'fakepython')
-    cmd = (sys.executable, fakepython, '-mpre_commit', *sys.argv[1:])
+    if sys.platform != 'win32':
+        exe = os.path.join(cache_dest, 'python')
+    else:
+        exe = os.path.join(cache_dest, 'python.exe')
+
+    cmd = (exe, '-mpre_commit', *sys.argv[1:])
     if sys.platform == 'win32':  # https://bugs.python.org/issue19124
         import subprocess
 
diff --git a/testing/zipapp/make b/testing/zipapp/make
@@ -2,6 +2,8 @@
 import argparse
 import base64
 import hashlib
+import importlib.resources
+import io
 import os.path
 import shutil
 import subprocess
@@ -30,11 +32,25 @@ def _check_no_shared_objects(wheeldir: str) -> None:
                     raise AssertionError(zip_filename, filename)
 
 
+def _add_shim(dest: str) -> None:
+    shim = os.path.join(HERE, 'python')
+    shutil.copy(shim, dest)
+
+    bio = io.BytesIO()
+    with zipfile.ZipFile(bio, 'w') as zipf:
+        zipf.write(shim, arcname='__main__.py')
+
+    with open(os.path.join(dest, 'python.exe'), 'wb') as f:
+        f.write(importlib.resources.read_binary('distlib', 't32.exe'))
+        f.write(b'#!py.exe -3\n')
+        f.write(bio.getvalue())
+
+
 def _write_cache_key(version: str, wheeldir: str, dest: str) -> None:
     cache_hash = hashlib.sha256(f'{version}\n'.encode())
     for filename in sorted(os.listdir(wheeldir)):
         cache_hash.update(f'{filename}\n'.encode())
-    with open(os.path.join(HERE, 'fakepython'), 'rb') as f:
+    with open(os.path.join(HERE, 'python'), 'rb') as f:
         cache_hash.update(f.read())
     with open(os.path.join(dest, 'CACHE_KEY'), 'wb') as f:
         f.write(base64.urlsafe_b64encode(cache_hash.digest()).rstrip(b'='))
@@ -62,11 +78,13 @@ def main() -> int:
         _msg('validating wheels...')
         _check_no_shared_objects(wheeldir)
 
-        _msg('adding fakepython / __main__.py...')
-        shutil.copy(os.path.join(HERE, 'fakepython'), tmpdir)
+        _msg('adding __main__.py...')
         mainfile = os.path.join(tmpdir, '__main__.py')
         shutil.copy(os.path.join(HERE, 'entry'), mainfile)
 
+        _msg('adding shim...')
+        _add_shim(tmpdir)
+
         _msg('copying file_lock.py...')
         file_lock_py = os.path.join(HERE, '../../pre_commit/file_lock.py')
         file_lock_py_dest = os.path.join(tmpdir, 'pre_commit/file_lock.py')
diff --git a/testing/zipapp/python b/testing/zipapp/python
@@ -5,7 +5,10 @@ import os.path
 import runpy
 import sys
 
-HERE = os.path.dirname(os.path.realpath(__file__))
+# an exe-zipapp will have a __file__ of shim.exe/__main__.py
+EXE = __file__ if os.path.isfile(__file__) else os.path.dirname(__file__)
+EXE = os.path.realpath(EXE)
+HERE = os.path.dirname(EXE)
 WHEELDIR = os.path.join(HERE, 'wheels')
 SITE_DIRS = frozenset(('dist-packages', 'site-packages'))
 
@@ -24,7 +27,7 @@ def main() -> int:
         for wheel in sorted(os.listdir(WHEELDIR)):
             sys.path.append(os.path.join(WHEELDIR, wheel))
         if args.m == 'pre_commit' or args.m.startswith('pre_commit.'):
-            sys.executable = os.path.abspath(__file__)
+            sys.executable = EXE
         sys.argv[1:] = rest
         runpy.run_module(args.m, run_name='__main__', alter_sys=True)
         return 0
