Ensure the correct source URL is stored with an install. (#260)

Fixes #259

Author: zooba

src/manage/indexutils.py

@@ -225,11 +225,11 @@ def find_to_install(self, tag, *, loose_company=False, prefer_prerelease=False):
         tag_list = [tag] if tag else []
         LOGGER.debug("Finding %s to install", tag_list)
         for i in self.find_all(tag_list, loose_company=loose_company, with_prerelease=prefer_prerelease):
-            return i
+            return {**i, "source": self.source_url}
         if not loose_company:
             return self.find_to_install(tag, loose_company=True, prefer_prerelease=prefer_prerelease)
         if not prefer_prerelease:
             for i in self.find_all(tag_list, loose_company=loose_company, with_prerelease=True):
-                return i
+                return {**i, "source": self.source_url}
         LOGGER.debug("No install found for %s", tag_list)
         raise LookupError(tag)

src/manage/install_command.py

@@ -528,6 +528,36 @@ def _restore_site(cmd, state):
             LOGGER.verbose("TRACEBACK", exc_info=True)
 
 
+def _sanitise_install(cmd, install):
+    """Prepares install metadata for storing locally.
+    
+    This includes:
+    * filtering out disabled shortcuts
+    * preserving original shortcuts
+    * sanitising URLs
+    """
+
+    if "shortcuts" in install:
+        # This saves our original set of shortcuts, so a later repair operation
+        # can enable those that were originally disabled.
+        shortcuts = install.setdefault("__original-shortcuts", install["shortcuts"])
+        if cmd.enable_shortcut_kinds or cmd.disable_shortcut_kinds:
+            orig_shortcuts = shortcuts
+            shortcuts = []
+            for s in orig_shortcuts:
+                if cmd.enable_shortcut_kinds and s["kind"] not in cmd.enable_shortcut_kinds:
+                    continue
+                if cmd.disable_shortcut_kinds and s["kind"] in cmd.disable_shortcut_kinds:
+                    continue
+                shortcuts.append(s)
+        install["shortcuts"] = shortcuts
+
+    install["url"] = sanitise_url(install["url"])
+    # If there's a non-empty and non-default source, sanitise it
+    if install.get("source") and install["source"] != cmd.fallback_source:
+        install["source"] = sanitise_url(install["source"])
+
+
 def _install_one(cmd, source, install, *, target=None):
     if cmd.repair:
         LOGGER.info("Repairing %s.", install['display-name'])
@@ -596,21 +626,7 @@ def _install_one(cmd, source, install, *, target=None):
             )
             raise
 
-        if "shortcuts" in install:
-            # This saves our original set of shortcuts, so a later repair operation
-            # can enable those that were originally disabled.
-            shortcuts = install.setdefault("__original-shortcuts", install["shortcuts"])
-            if cmd.enable_shortcut_kinds:
-                shortcuts = [s for s in shortcuts
-                             if s["kind"] in cmd.enable_shortcut_kinds]
-            if cmd.disable_shortcut_kinds:
-                shortcuts = [s for s in shortcuts
-                             if s["kind"] not in cmd.disable_shortcut_kinds]
-            install["shortcuts"] = shortcuts
-
-        install["url"] = sanitise_url(install["url"])
-        if source != cmd.fallback_source:
-            install["source"] = sanitise_url(source)
+        _sanitise_install(cmd, install)
 
         LOGGER.debug("Write __install__.json to %s", dest)
         with open(dest / "__install__.json", "w", encoding="utf-8") as f:
@@ -850,10 +866,12 @@ def execute(cmd):
                 if cmd.download:
                     LOGGER.info("Downloading %s", install["display-name"])
                     package = _download_one(cmd, source, install, cmd.download, must_copy=True)
-                    download_index["versions"].append({
+                    install_idx = {
                         **install,
                         "url": package.name,
-                    })
+                    }
+                    install_idx.pop("source", None)
+                    download_index["versions"].append(install_idx)
                 else:
                     _install_one(cmd, source, install)
         except ArgumentError:

tests/test_install_command.py

@@ -328,3 +328,88 @@ def test_install_from_script(tmp_path, assert_log):
         assert_log.skip_until("Installing %s", ["Test 1.1 (32)"]),
         ("Tag: %s\\\\%s", ["Test", "1.1-32"]),
     )
+
+
+def test_sanitise_install_urls():
+    class Cmd:
+        enable_shortcut_kinds = []
+        disable_shortcut_kinds = []
+        fallback_source = None
+
+    i = {
+        "url": "http://user:placeholder@example.com/package.zip",
+        "source": "http://user:placeholder@example.com/index.json",
+    }
+
+    IC._sanitise_install(Cmd, i)
+
+    assert i["url"] == "http://example.com/package.zip"
+    assert i["source"] == "http://example.com/index.json"
+
+
+def test_sanitise_install_fallback_urls():
+    class Cmd:
+        enable_shortcut_kinds = []
+        disable_shortcut_kinds = []
+        fallback_source = "http://user:placeholder@example.com/index.json"
+
+    i = {
+        "url": "http://user:placeholder@example.com/package.zip",
+        "source": "http://user:placeholder@example.com/index.json",
+    }
+
+    IC._sanitise_install(Cmd, i)
+
+    assert i["url"] == "http://example.com/package.zip"
+    assert i["source"] == "http://user:placeholder@example.com/index.json"
+
+
+def test_sanitise_install_shortcuts():
+    class Cmd:
+        enable_shortcut_kinds = []
+        disable_shortcut_kinds = []
+        fallback_source = None
+
+    i = {
+        "url": "",
+        "shortcuts": [dict(kind=a) for a in "abc"],
+    }
+
+    IC._sanitise_install(Cmd, i)
+
+    assert [j["kind"] for j in i["shortcuts"]] == ["a", "b", "c"]
+    assert [j["kind"] for j in i["__original-shortcuts"]] == ["a", "b", "c"]
+
+
+def test_sanitise_install_shortcuts_disable():
+    class Cmd:
+        enable_shortcut_kinds = []
+        disable_shortcut_kinds = ["b"]
+        fallback_source = None
+
+    i = {
+        "url": "",
+        "shortcuts": [dict(kind=a) for a in "abc"],
+    }
+
+    IC._sanitise_install(Cmd, i)
+
+    assert [j["kind"] for j in i["shortcuts"]] == ["a", "c"]
+    assert [j["kind"] for j in i["__original-shortcuts"]] == ["a", "b", "c"]
+
+
+def test_sanitise_install_shortcuts_enable():
+    class Cmd:
+        enable_shortcut_kinds = ["b"]
+        disable_shortcut_kinds = []
+        fallback_source = None
+
+    i = {
+        "url": "",
+        "shortcuts": [dict(kind=a) for a in "abc"],
+    }
+
+    IC._sanitise_install(Cmd, i)
+
+    assert [j["kind"] for j in i["shortcuts"]] == ["b"]
+    assert [j["kind"] for j in i["__original-shortcuts"]] == ["a", "b", "c"]