Adds shebang_templates configuration for overriding shebangs. (#359)

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,61 @@ def _find_on_path(cmd, full_cmd):
     }
 
 
+def _replace_templates(cmd, line, windowed):
+    # Override can be the entire line or just the first argument
+    # Override can be the entire line (including args) or just the first argument
+    m = re.match(r"^#!\s*([^\s]+)(.*)$", line)
+
+    if not m:
+        return None, None
+
+    full_key = (m.group(1) + m.group(2)).strip()
+    if full_key in cmd.shebang_templates:
+        template_key = full_key
+        suffix = ""
+    elif m.group(1) in cmd.shebang_templates:
+        template_key = m.group(1)
+        suffix = m.group(2)
+    else:
+        return None, None
+
+    new_cmd = cmd.shebang_templates[template_key]
+    LOGGER.verbose("Using '%s' from configuration file in place of shebang '%s'",
+                   new_cmd, template_key)
+    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}{suffix}"
+    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 +205,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 +235,20 @@ 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)
+                raise
         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": "pyw",
+        "/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_install_command.py

@@ -244,6 +244,7 @@ def __init__(self, tmp_path, *args, **kwargs):
         self.repair = kwargs.pop("repair", False)
         self.shebang_can_run_anything = kwargs.pop("shebang_can_run_anything", False)
         self.shebang_can_run_anything_silently = kwargs.pop("shebang_can_run_anything_silently", False)
+        self.shebang_templates = {}
         self.source = kwargs.pop("source", "http://example.com/index.json")
         self.target = kwargs.pop("target", None)
         if self.target:

tests/test_scriptutils.py

@@ -6,10 +6,13 @@
 
 from pathlib import PurePath
 
+import manage.scriptutils as SU
 from manage.scriptutils import (
     find_install_from_script,
     _find_shebang_command,
+    _parse_shebang,
     _read_script,
+    _replace_templates,
     NewEncoding,
     _maybe_quote,
     quote_args,
@@ -274,3 +277,61 @@ 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/python -Es", "Test1", None),
+    ("#! /usr/bin/pythonw", "Test1", None),
+    ("#!  /usr/bin/python2", "Test2", None),
+    ("#!   /usr/bin/pythonw2", "Test2", None),
+    ("#!  /usr/bin/python3", "PythonCore3", None),
+    ("#!   /usr/bin/pythonw3", "PythonCore3", None),
+    ("#!  custom", None, "#!CUSTOM"),
+    ("#! full line custom", None, "#!CUSTOM2"),
+    ("#!full line custom with extra", None, None),
+    ("custom", None, None),
+    ("full line custom", None, None),
+]])
+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"),
+        dict(id="Test3", company="Test", tag="3.2"),
+        dict(id="PythonCore3", company="PythonCore", tag="3.2"),
+    ]
+    fake_config.shebang_templates = {
+        "/usr/bin/python": "py",
+        "/usr/bin/pythonw": "pyw",
+        "/usr/bin/python2": "py -V:Test/2",
+        "/usr/bin/pythonw2": "pyw -V:Test/2",
+        "/usr/bin/python3": "py -3.2",
+        "/usr/bin/pythonw3": "pyw -3.2",
+        "custom": "CUSTOM",
+        "full line custom": "CUSTOM2",
+    }
+    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:
+        assert not actual
+        assert not actual_line
+
+
+def test_parse_shebang_templates(monkeypatch):
+    class Cmd:
+        shebang_templates = True
+
+    expect = {"an": "install"}
+    monkeypatch.setattr(SU, "_replace_templates", lambda *a: (expect, None))
+    actual = _parse_shebang(Cmd, "Anything at all")
+    assert expect == actual
+
+    expect = {"id": "COMMAND"}
+    monkeypatch.setattr(SU, "_replace_templates", lambda *a: (None, "#!COMMAND"))
+    monkeypatch.setattr(SU, "_find_shebang_command", lambda cmd, full_cmd, **kw: {"id": full_cmd})
+    actual = _parse_shebang(Cmd, "Anything at all", windowed=False)
+    assert expect == actual