Improve handling when install fails. (#323)

zooba · web-flow · commit 07c4712c21c6 · 2026-04-24T11:02:54.000+01:00
Fixes #322
diff --git a/src/manage/install_command.py b/src/manage/install_command.py
@@ -108,7 +108,7 @@ def download_package(cmd, install, dest, cache, *, on_progress=None, urlopen=_ur
                 LOGGER.verbose("Using bundled file at %s", bundled)
                 return bundled
 
-    unlink(dest, "Removing old download is taking some time. " + 
+    unlink(dest, "Removing old download is taking some time. " +
                  "Please continue to wait, or press Ctrl+C to abort.")
 
     def _find_creds(url):
@@ -528,36 +528,6 @@ 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'])
@@ -574,46 +544,87 @@ def _install_one(cmd, source, install, *, target=None):
     package = _download_one(cmd, source, install, cmd.download_dir)
 
     dest = target or (cmd.install_dir / install["id"])
+    metadata_dest = dest / "__install__.json"
 
     preserved_site = _preserve_site(cmd, dest, install)
 
     LOGGER.verbose("Extracting %s to %s", package, dest)
-    if not cmd.repair:
-        try:
-            rmtree(
-                dest,
-                "Removing the previous install is taking some time. " +
-                "Ensure Python is not running, and continue to wait " +
-                "or press Ctrl+C to abort.",
-                remove_ext_first=("exe", "dll", "json"),
-            )
-        except FileExistsError:
-            LOGGER.error(
-                "Unable to remove previous install. " +
-                "Please check your packages directory at %s for issues.",
-                dest.parent
-            )
-            raise
-        except FilesInUseError:
-            LOGGER.error(
-                "Unable to remove previous install because files are still in use. " +
-                "Please ensure Python is not currently running."
+    try:
+        if not cmd.repair:
+            _remove_existing(dest)
+
+        with ProgressPrinter("Extracting", maxwidth=CONSOLE_MAX_WIDTH) as on_progress:
+            extract_package(package, dest, on_progress=on_progress, repair=cmd.repair)
+
+        if target:
+            unlink(
+                metadata_dest,
+                "Removing metadata from the install is taking some time. Please " +
+                "continue to wait, or press Ctrl+C to abort."
             )
-            raise
+        else:
+            _finalize_metadata(cmd, install, metadata_dest)
 
-    with ProgressPrinter("Extracting", maxwidth=CONSOLE_MAX_WIDTH) as on_progress:
-        extract_package(package, dest, on_progress=on_progress, repair=cmd.repair)
+            LOGGER.debug("Write __install__.json to %s", dest)
+            with open(metadata_dest, "w", encoding="utf-8") as f:
+                json.dump(install, f, default=str)
 
-    if target:
-        unlink(
-            dest / "__install__.json",
-            "Removing metadata from the install is taking some time. Please " +
-            "continue to wait, or press Ctrl+C to abort."
+    finally:
+        # May be letting an exception bubble out here, so we'll handle and log
+        # here rather than letting any new ones leave.
+        try:
+            if dest.is_dir():
+                # Install may be broken at this point, but we'll put site back anyway
+                _restore_site(cmd, preserved_site)
+            else:
+                # Install is certainly broken, but we don't want to delete user files
+                # Just warn, until we come up with a better idea
+                LOGGER.warn("This runtime has been lost due to an error, you will "
+                            "need to reinstall.")
+        except Exception:
+            LOGGER.warn("Unexpected failure finalizing install. See log file for details")
+            LOGGER.verbose("TRACEBACK", exc_info=True)
+
+    LOGGER.verbose("Install complete")
+
+
+def _remove_existing(install_dir):
+    try:
+        rmtree(
+            install_dir,
+            "Removing the previous install is taking some time. " +
+            "Ensure Python is not running, and continue to wait " +
+            "or press Ctrl+C to abort.",
+            remove_ext_first=("exe", "dll", "json"),
         )
-    else:
+    except FileExistsError:
+        LOGGER.error(
+            "Unable to remove previous install. " +
+            "Please check your packages directory at %s for issues.",
+            install_dir.parent
+        )
+        raise
+    except FilesInUseError:
+        LOGGER.error(
+            "Unable to remove previous install because files are still in use. " +
+            "Please ensure Python is not currently running."
+        )
+        raise
+
+
+def _finalize_metadata(cmd, install, merge_from=None):
+    """Prepares install metadata for storing locally.
+
+    This includes:
+    * filtering out disabled shortcuts
+    * preserving original shortcuts
+    * sanitising URLs
+    """
+
+    if merge_from:
         try:
-            with open(dest / "__install__.json", "r", encoding="utf-8-sig") as f:
-                LOGGER.debug("Updating from __install__.json in %s", dest)
+            with open(merge_from, "r", encoding="utf-8-sig") as f:
+                LOGGER.debug("Updating from __install__.json in %s", merge_from.parent)
                 for k, v in json.load(f).items():
                     if not install.setdefault(k, v):
                         install[k] = v
@@ -626,15 +637,25 @@ def _install_one(cmd, source, install, *, target=None):
             )
             raise
 
-        _sanitise_install(cmd, install)
-
-        LOGGER.debug("Write __install__.json to %s", dest)
-        with open(dest / "__install__.json", "w", encoding="utf-8") as f:
-            json.dump(install, f, default=str)
-
-    _restore_site(cmd, preserved_site)
+    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
 
-    LOGGER.verbose("Install complete")
+    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 _merge_existing_index(versions, index_json):
diff --git a/tests/test_install_command.py b/tests/test_install_command.py
@@ -233,6 +233,7 @@ def __init__(self, tmp_path, *args, **kwargs):
         self.download = kwargs.get("download")
         if self.download:
             self.download = tmp_path / self.download
+        self.download_dir = tmp_path / kwargs.get("download_dir", "_cache")
         self.dry_run = kwargs.get("dry_run", True)
         self.fallback_source = kwargs.get("fallback_source")
         self.force = kwargs.get("force", True)
@@ -330,7 +331,85 @@ def test_install_from_script(tmp_path, assert_log):
     )
 
 
-def test_sanitise_install_urls():
+def test_failed_install_unwind(tmp_path, monkeypatch, assert_log):
+    cmd = InstallCommandTestCmd(tmp_path, "1.0", force=False)
+    cmd.dry_run = False
+    cmd.preserve_site_on_upgrade = True
+
+    inst = cmd.installs[0]
+    inst.setdefault("shortcuts", []).append({
+        "kind": "site-dirs", "dirs": ["test-site"],
+    })
+
+    target = tmp_path / "target"
+    test_file = target / "test-site/file.txt"
+    test_file.parent.mkdir(parents=True, exist_ok=True)
+    test_file.write_text("Before")
+
+    def remove_existing(*args):
+        pass
+
+    def download_one(*args, **kwargs):
+        return "<package>"
+
+    def extract_package(package, dest, *args, **kwargs):
+        # site dir should be gone
+        assert not test_file.is_file()
+        # create the target directory
+        dest.mkdir(parents=True, exist_ok=True)
+        # interrupt the install process
+        raise RuntimeError("Failed to extract for test reasons")
+
+    monkeypatch.setattr(IC, "_remove_existing", remove_existing)
+    monkeypatch.setattr(IC, "_download_one", download_one)
+    monkeypatch.setattr(IC, "extract_package", extract_package)
+
+    with pytest.raises(RuntimeError):
+        IC._install_one(cmd, "<source>", inst, target=target)
+
+    # site dir should be back
+    assert test_file.is_file()
+    assert test_file.read_text() == "Before"
+
+
+def test_failed_install_unwind_dont_clobber(tmp_path, monkeypatch, assert_log):
+    cmd = InstallCommandTestCmd(tmp_path, "1.0", force=False)
+    cmd.dry_run = False
+    cmd.preserve_site_on_upgrade = True
+
+    inst = cmd.installs[0]
+    inst.setdefault("shortcuts", []).append({
+        "kind": "site-dirs", "dirs": ["test-site"],
+    })
+
+    target = tmp_path / "test-site"
+    test_file = target / "file.txt"
+    test_file.parent.mkdir(parents=True, exist_ok=True)
+    test_file.write_text("Before")
+
+    def download_one(*args, **kwargs):
+        return "<package>"
+
+    def extract_package(package, dest, *args, **kwargs):
+        # site dir should be gone
+        assert not test_file.is_file()
+        # create a new one - it should be preserved
+        test_file.parent.mkdir(parents=True, exist_ok=True)
+        test_file.write_text("After")
+        # interrupt the install process
+        raise RuntimeError("Failed to extract for test reasons")
+
+    monkeypatch.setattr(IC, "_download_one", download_one)
+    monkeypatch.setattr(IC, "extract_package", extract_package)
+
+    with pytest.raises(RuntimeError):
+        IC._install_one(cmd, "<source>", inst, target=target)
+
+    # Ensure file we extracted is still there
+    assert test_file.read_text() == "After"
+
+
+def test_finalize_metadata_urls():
     class Cmd:
         enable_shortcut_kinds = []
         disable_shortcut_kinds = []
@@ -341,13 +420,13 @@ class Cmd:
         "source": "http://user:placeholder@example.com/index.json",
     }
 
-    IC._sanitise_install(Cmd, i)
+    IC._finalize_metadata(Cmd, i)
 
     assert i["url"] == "http://example.com/package.zip"
     assert i["source"] == "http://example.com/index.json"
 
 
-def test_sanitise_install_fallback_urls():
+def test_finalize_metadata_fallback_urls():
     class Cmd:
         enable_shortcut_kinds = []
         disable_shortcut_kinds = []
@@ -358,13 +437,13 @@ class Cmd:
         "source": "http://user:placeholder@example.com/index.json",
     }
 
-    IC._sanitise_install(Cmd, i)
+    IC._finalize_metadata(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():
+def test_finalize_metadata_shortcuts():
     class Cmd:
         enable_shortcut_kinds = []
         disable_shortcut_kinds = []
@@ -375,13 +454,13 @@ class Cmd:
         "shortcuts": [dict(kind=a) for a in "abc"],
     }
 
-    IC._sanitise_install(Cmd, i)
+    IC._finalize_metadata(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():
+def test_finalize_metadata_shortcuts_disable():
     class Cmd:
         enable_shortcut_kinds = []
         disable_shortcut_kinds = ["b"]
@@ -392,13 +471,13 @@ class Cmd:
         "shortcuts": [dict(kind=a) for a in "abc"],
     }
 
-    IC._sanitise_install(Cmd, i)
+    IC._finalize_metadata(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():
+def test_finalize_metadata_shortcuts_enable():
     class Cmd:
         enable_shortcut_kinds = ["b"]
         disable_shortcut_kinds = []
@@ -409,7 +488,32 @@ class Cmd:
         "shortcuts": [dict(kind=a) for a in "abc"],
     }
 
-    IC._sanitise_install(Cmd, i)
+    IC._finalize_metadata(Cmd, i)
 
     assert [j["kind"] for j in i["shortcuts"]] == ["b"]
     assert [j["kind"] for j in i["__original-shortcuts"]] == ["a", "b", "c"]
+
+
+def test_finalize_metadata_merge_from(tmp_path):
+    class Cmd:
+        enable_shortcut_kinds = []
+        disable_shortcut_kinds = []
+        fallback_source = None
+
+    merge_from = tmp_path / "file.json"
+    test_url_1 = "https://example.com/"
+    test_url_2 = "https://example.com/path2"
+
+    # merge_from does not exist, but we shouldn't fail
+    i = {"url": test_url_1}
+    IC._finalize_metadata(Cmd, i, merge_from)
+    assert i["url"] == test_url_1
+
+    # Update missing fields from merge_from, but don't touch existing ones
+    with open(merge_from, "w", encoding="utf-8") as f:
+        json.dump({"url": test_url_1, "data1": "b", "data2": "c"}, f)
+    i = {"url": test_url_2, "data1": "a"}
+    IC._finalize_metadata(Cmd, i, merge_from)
+    assert i["url"] == test_url_2
+    assert i["data1"] == "a"
+    assert i["data2"] == "c"
