Merge pull request #893 from pre-commit/shebang

asottile · web-flow · commit 1d40cc210462 · 2018-12-25T21:05:21.000-08:00
Pick a better python shebang for hook executable
diff --git a/pre_commit/commands/install_uninstall.py b/pre_commit/commands/install_uninstall.py
@@ -8,6 +8,7 @@
 
 from pre_commit import git
 from pre_commit import output
+from pre_commit.languages import python
 from pre_commit.repository import repositories
 from pre_commit.util import cmd_output
 from pre_commit.util import make_executable
@@ -43,6 +44,16 @@ def is_our_script(filename):
     return any(h in contents for h in (CURRENT_HASH,) + PRIOR_HASHES)
 
 
+def shebang():
+    if sys.platform == 'win32':
+        py = 'python'
+    else:
+        py = python.get_default_version()
+        if py == 'default':
+            py = 'python'
+    return '#!/usr/bin/env {}'.format(py)
+
+
 def install(
         runner, store, overwrite=False, hooks=False, hook_type='pre-commit',
         skip_on_missing_conf=False,
@@ -84,6 +95,8 @@ def install(
         before, rest = contents.split(TEMPLATE_START)
         to_template, after = rest.split(TEMPLATE_END)
 
+        before = before.replace('#!/usr/bin/env python', shebang())
+
         hook_file.write(before + TEMPLATE_START)
         for line in to_template.splitlines():
             var = line.split()[0]
diff --git a/tests/commands/install_uninstall_test.py b/tests/commands/install_uninstall_test.py
@@ -17,7 +17,9 @@
 from pre_commit.commands.install_uninstall import install_hooks
 from pre_commit.commands.install_uninstall import is_our_script
 from pre_commit.commands.install_uninstall import PRIOR_HASHES
+from pre_commit.commands.install_uninstall import shebang
 from pre_commit.commands.install_uninstall import uninstall
+from pre_commit.languages import python
 from pre_commit.runner import Runner
 from pre_commit.util import cmd_output
 from pre_commit.util import make_executable
@@ -45,6 +47,24 @@ def test_is_previous_pre_commit(tmpdir):
     assert is_our_script(f.strpath)
 
 
+def test_shebang_windows():
+    with mock.patch.object(sys, 'platform', 'win32'):
+        assert shebang() == '#!/usr/bin/env python'
+
+
+def test_shebang_otherwise():
+    with mock.patch.object(sys, 'platform', 'posix'):
+        assert 'default' not in shebang()
+
+
+def test_shebang_returns_default():
+    with mock.patch.object(sys, 'platform', 'posix'):
+        with mock.patch.object(
+            python, 'get_default_version', return_value='default',
+        ):
+            assert shebang() == '#!/usr/bin/env python'
+
+
 def test_install_pre_commit(tempdir_factory, store):
     path = git_dir(tempdir_factory)
     runner = Runner(path, C.CONFIG_FILE)
