From 36670b49a0339168145aa4ed53bb56f20d8f6a18 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Tue, 3 Aug 2021 17:47:47 +0530
Subject: [PATCH] Make creation of kitty-themes.zip atomic

---
 kittens/themes/collection.py | 27 ++++++++++++++++++++++-----
 kittens/themes/main.py       |  4 +---
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/kittens/themes/collection.py b/kittens/themes/collection.py
index 06593360e..df763383d 100644
--- a/kittens/themes/collection.py
+++ b/kittens/themes/collection.py
@@ -8,9 +8,10 @@ import json
 import os
 import re
 import shutil
+import tempfile
 import zipfile
 from contextlib import suppress
-from typing import Any, Callable, Dict, Match, Optional
+from typing import Any, Callable, Dict, Iterator, Match, Optional
 from urllib.error import HTTPError
 from urllib.request import Request, urlopen
 
@@ -48,10 +49,20 @@ def fetch_themes(
             return dest_path
         raise
     m.etag = res.headers.get('etag') or ''
-    with open(dest_path, 'wb') as f:
-        shutil.copyfileobj(res, f)
-    with zipfile.ZipFile(dest_path, 'a') as zf:
-        zf.comment = json.dumps({'etag': m.etag, 'timestamp': m.timestamp.isoformat()}).encode('utf-8')
+
+    needs_delete = False
+    try:
+        with tempfile.NamedTemporaryFile(suffix='-' + os.path.basename(dest_path), dir=os.path.dirname(dest_path), delete=False) as f:
+            needs_delete = True
+            shutil.copyfileobj(res, f)
+            f.flush()
+            with zipfile.ZipFile(f.name, 'a') as zf:
+                zf.comment = json.dumps({'etag': m.etag, 'timestamp': m.timestamp.isoformat()}).encode('utf-8')
+            os.replace(f.name, dest_path)
+            needs_delete = False
+    finally:
+        if needs_delete:
+            os.unlink(f.name)
     return dest_path
 
 
@@ -171,6 +182,12 @@ class Themes:
     def __init__(self) -> None:
         self.themes: Dict[str, Theme] = {}
 
+    def __len__(self) -> int:
+        return len(self.themes)
+
+    def __iter__(self) -> Iterator[Theme]:
+        return iter(self.themes.values())
+
     def load_from_zip(self, path_to_zip: str) -> None:
         with zipfile.ZipFile(path_to_zip, 'r') as zf:
             for name in zf.namelist():
diff --git a/kittens/themes/main.py b/kittens/themes/main.py
index 5f23ee912..7e2322bd0 100644
--- a/kittens/themes/main.py
+++ b/kittens/themes/main.py
@@ -58,13 +58,11 @@ class ThemesHandler(Handler):
                 self.quit_loop(1)
                 return
             else:
-                self.themes: Themes = themes_or_exception
+                self.all_themes: Themes = themes_or_exception
                 self.state = State.browsing
             self.draw_screen()
 
         def fetch() -> None:
-            import time
-            time.sleep(15)
             try:
                 themes: Union[Themes, str] = load_themes()
             except Exception: