Adds shebang_templates configuration for overriding shebangs.

Fixes #307
Fixes #348

Author: zooba

src/manage/commands.py

@@ -242,6 +242,9 @@ def execute(self):
     "include_unmanaged": (config_bool, None, "env"),
     "shebang_can_run_anything": (config_bool, None, "env"),
     "shebang_can_run_anything_silently": (config_bool, None, "env"),
+    # Mapping from shebang template to '-V:Company/Tag' argument or an
+    # executable path. The latter requires 'shebang_can_run_anything'.
+    "shebang_templates": (dict, config_dict_merge),
     # Typically configured to '%VIRTUAL_ENV%' to pick up the active environment
     "virtual_env": (str, None, "env", "path"),
 
@@ -347,6 +350,7 @@ class BaseCommand:
     virtual_env = None
     shebang_can_run_anything = True
     shebang_can_run_anything_silently = False
+    shebang_templates = {}
     welcome_on_update = False
 
     log_file = None
@@ -366,7 +370,7 @@ class BaseCommand:
     launcher_exe = None
     launcherw_exe = None
 
-    source_settings = None
+    source_settings = {}
 
     show_help = False
 

src/manage/scriptutils.py

@@ -121,7 +121,50 @@ def _find_on_path(cmd, full_cmd):
     }
 
 
+def _replace_templates(cmd, line, windowed):
+    # Override can be the entire line or just the first argument
+    shebang = re.match(r"#!\s*(.+)(.*)", line) or re.match(r"#!\s*([^\s]+)(.*)", line)
+
+    if not shebang or shebang.group(1) not in cmd.shebang_templates:
+        return None, None
+
+    new_cmd = cmd.shebang_templates[shebang.group(1)]
+    LOGGER.verbose("Using '%s' from configuration file in place of shebang '%s'",
+                   new_cmd, shebang.group(1))
+    install = None
+    if new_cmd.startswith("py -V:"):
+        install = cmd.get_install_to_run(new_cmd[6:], windowed=windowed)
+    elif new_cmd.startswith("pyw -V:"):
+        install = cmd.get_install_to_run(new_cmd[7:], windowed=True)
+    elif new_cmd.startswith("py -3"):
+        install = cmd.get_install_to_run(f"PythonCore/{new_cmd[4:]}", windowed=windowed)
+    elif new_cmd.startswith("pyw -3"):
+        install = cmd.get_install_to_run(f"PythonCore/{new_cmd[5:]}", windowed=True)
+    elif new_cmd == "py":
+        install = cmd.get_install_to_run(windowed=windowed)
+    elif new_cmd == "pyw":
+        install = cmd.get_install_to_run(windowed=True)
+    else:
+        # Recreate the shebang with the alternate command and continue.
+        line = f"#!{new_cmd}{shebang.group(2)}"
+    return install, line
+
+
 def _parse_shebang(cmd, line, *, windowed=None):
+    # To silence our warning when we get the path from config file
+    run_anything_silently = False
+
+    # First check the user-provided overrides
+    if cmd.shebang_templates:
+        install, new_line = _replace_templates(cmd, line, windowed)
+        if install:
+            return install
+        if new_line:
+            # We don't warn about custom executables if they've come from
+            # the config file, unless they don't exist or are disabled.
+            run_anything_silently = True
+            line = new_line
+
     # For /usr[/local]/bin, we look for a matching alias name.
     shebang = re.match(r"#!\s*/usr/(?:local/)?bin/(?!env\b)([^\\/\s]+).*", line)
     if shebang:
@@ -151,7 +194,7 @@ def _parse_shebang(cmd, line, *, windowed=None):
         # If not, warn and do regular PATH search
         if cmd.shebang_can_run_anything or cmd.shebang_can_run_anything_silently:
             i = _find_on_path(cmd, full_cmd)
-            if not cmd.shebang_can_run_anything_silently:
+            if not cmd.shebang_can_run_anything_silently and not run_anything_silently:
                 LOGGER.warn("A shebang '%s' was found but could not be matched "
                             "to an installed runtime, so it will be treated as "
                             "an arbitrary command.", full_cmd)
@@ -181,14 +224,19 @@ def _parse_shebang(cmd, line, *, windowed=None):
         except LookupError:
             pass
         if cmd.shebang_can_run_anything or cmd.shebang_can_run_anything_silently:
-            if not cmd.shebang_can_run_anything_silently:
+            if not cmd.shebang_can_run_anything_silently and not run_anything_silently:
                 LOGGER.warn("A shebang '%s' was found but does not match any "
-                            "supported template (e.g. '/usr/bin/python'), so it "
-                            "will be treated as an arbitrary command.", full_cmd)
+                            "supported or configured template (e.g. "
+                            "'/usr/bin/python'), so it will be treated as an "
+                            "arbitrary command.", full_cmd)
                 LOGGER.warn("To prevent execution of programs that are not "
                             "Python runtimes, set 'shebang_can_run_anything' to "
                             "'false' in your configuration file.")
-            return _find_on_path(cmd, full_cmd)
+            try:
+                return _find_on_path(cmd, full_cmd)
+            except LookupError as ex:
+                LOGGER.error("Could not launch '%s'. Using default interpreter "
+                             "instead.", full_cmd)
         else:
             LOGGER.warn("A shebang '%s' was found, but could not be matched "
                         "to an installed runtime.", full_cmd)

src/pymanager.json

@@ -34,6 +34,17 @@
     "launcherw_exe": "./templates/launcherw.exe",
     "welcome_on_update": true,
 
+    "shebang_templates": {
+        "/usr/bin/python": "py",
+        "/usr/bin/pythonw": "py",
+        "/usr/bin/python3": "py",
+        "/usr/bin/pythonw3": "pyw",
+        "/usr/local/bin/python": "py",
+        "/usr/local/bin/pythonw": "pyw",
+        "/usr/local/bin/python3": "py",
+        "/usr/local/bin/pythonw3": "pyw"
+    },
+
     "source_settings": {
         "https://www.python.org/ftp/python/index-windows.json": {
             "requires_signature": true,

tests/conftest.py

@@ -162,28 +162,32 @@ def __init__(self, root, installs=[]):
         self.installs = list(installs)
         self.shebang_can_run_anything = True
         self.shebang_can_run_anything_silently = False
+        self.shebang_templates = {}
         self.scratch = {}
 
     def get_installs(self, *, include_unmanaged=False, set_default=True):
         if include_unmanaged:
             return self.installs
         return [i for i in self.installs if not i.get("unmanaged", 0)]
 
-    def get_install_to_run(self, tag, *, windowed=False):
+    def get_install_to_run(self, tag=None, script=None, *, windowed=False):
         if windowed:
             i = self.get_install_to_run(tag)
             target = [t for t in i.get("run-for", []) if t.get("windowed")]
             if target:
                 return {**i, "executable": i["prefix"] / target[0]["target"]}
             return i
 
-        company, _, tag = tag.replace("/", "\\").rpartition("\\")
-        try:
-            found = [i for i in self.installs
-                     if (not tag or i["tag"] == tag) and (not company or i["company"] == company)]
-        except LookupError as ex:
-            # LookupError is expected from this function, so make sure we don't raise it here
-            raise RuntimeError from ex
+        if not tag:
+            found = [i for i in self.installs if i.get("default")]
+        else:
+            company, _, tag = tag.replace("/", "\\").rpartition("\\")
+            try:
+                found = [i for i in self.installs
+                         if (not tag or i["tag"] == tag) and (not company or i["company"] == company)]
+            except LookupError as ex:
+                # LookupError is expected from this function, so make sure we don't raise it here
+                raise RuntimeError from ex
         if found:
             return found[0]
         raise LookupError(tag)

tests/test_scriptutils.py

@@ -10,6 +10,7 @@
     find_install_from_script,
     _find_shebang_command,
     _read_script,
+    _replace_templates,
     NewEncoding,
     _maybe_quote,
     quote_args,
@@ -274,3 +275,27 @@ def test_quote_args(args, expect):
     assert expect == quote_args(args)
     # Test that our split function produces the same result
     assert args == split_args(expect), expect
+
+
+@pytest.mark.parametrize("line, expect_id, expect_line", [pytest.param(*a, id=a[0]) for a in [
+    ("#!/usr/bin/python", "Test1", None),
+    ("#!/usr/bin/python2", "Test2", None),
+    # TODO: More test cases
+]])
+def test_shebang_templates(fake_config, line, expect_id, expect_line):
+    fake_config.installs = [
+        dict(id="Test1", company="Test", tag="1", default=True),
+        dict(id="Test2", company="Test", tag="2"),
+    ]
+    fake_config.shebang_templates = {
+        "/usr/bin/python": "py",
+        "/usr/bin/python2": "py -V:Test/2",
+    }
+    actual, actual_line = _replace_templates(fake_config, line, False)
+    if expect_id:
+        assert actual
+        assert expect_id == actual["id"]
+    elif expect_line:
+        assert expect_line == actual_line
+    else:
+        pytest.fail("Invalid test")