fix: last_screenshot corrupt

fix #235

auto_derby/filetools.py

@@ -0,0 +1,34 @@
+# -*- coding=UTF-8 -*-
+# pyright: strict
+
+from __future__ import annotations
+
+import contextlib
+import os
+from typing import Text
+
+
+def force_rename(src: Text, dst: Text):
+    try:
+        os.rename(src, dst)
+    except FileExistsError:
+        os.unlink(dst)
+        os.rename(src, dst)
+
+
+@contextlib.contextmanager
+def atomic_save_path(
+    path: Text,
+    *,
+    temp_suffix: Text = ".tmp",
+    backup_suffix: Text = "",
+):
+    tmp_path = path + temp_suffix
+    yield tmp_path
+    if backup_suffix:
+        backup_path = path + backup_suffix
+        try:
+            force_rename(path, backup_path)
+        except FileNotFoundError:
+            pass
+    force_rename(tmp_path, path)

auto_derby/single_mode/race/race_result.py

@@ -12,6 +12,7 @@
 from ..context import Context
 from .globals import g
 from .race import Race
+from ... import filetools
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -129,20 +130,15 @@ def iterate() -> Iterator[RaceResult]:
 
 
 def prune(time_lt: datetime.datetime) -> None:
-    p = g.result_path
-    tmp_path = p + ".tmp"
-    with open(tmp_path, "w", encoding="utf-8") as f:
+    with filetools.atomic_save_path(
+        g.result_path,
+        backup_suffix="~",
+    ) as p, open(p, "w", encoding="utf-8") as f:
         for res in iterate():
             if res.time < time_lt:
                 continue
             json.dump(res.to_dict(), f, ensure_ascii=False)
             f.write("\n")
-    try:
-        os.unlink(p + "~")
-    except FileNotFoundError:
-        pass
-    os.rename(p, p + "~")
-    os.rename(tmp_path, p)
 
 
 def iterate_current(ctx: Context) -> Iterator[RaceResult]:

auto_derby/template.py

@@ -13,7 +13,7 @@
 from PIL.Image import Image
 from PIL.Image import open as open_image
 
-from . import clients, imagetools, mathtools
+from . import clients, imagetools, mathtools, filetools
 
 LOGGER = logging.getLogger(__name__)
 
@@ -40,7 +40,10 @@ def screenshot(*, max_age: float = 1) -> Image:
         g.screenshot_width = new_img.width
         new_img = new_img.convert("RGB")
         if g.last_screenshot_save_path:
-            new_img.save(g.last_screenshot_save_path)
+            with filetools.atomic_save_path(
+                g.last_screenshot_save_path,
+            ) as p:
+                new_img.save(p, format="PNG")
         LOGGER.debug("screenshot")
         _g.cached_screenshot = (dt.datetime.now(), new_img)
     return _g.cached_screenshot[1]