Wrap stored object concept into context manager

geo_activity_playground/core/similarity.py

@@ -10,6 +10,7 @@
 
 from .activities import ActivityRepository
 from .coordinates import get_distance
+from geo_activity_playground.core.tasks import stored_object
 
 
 fingerprint_path = pathlib.Path("Cache/activity_fingerprints.pickle")
@@ -25,40 +26,30 @@ def add_distance(distances, this, other, distance) -> None:
 
 
 def precompute_activity_distances(repository: ActivityRepository) -> None:
-    fingerprints: dict[int, int] = {}
-    if fingerprint_path.exists():
-        with open(fingerprint_path, "rb") as f:
-            fingerprints = pickle.load(f)
-
-    activity_ids = repository.activity_ids
-
-    activity_ids_without_fingerprint = [
-        activity_id for activity_id in activity_ids if activity_id not in fingerprints
-    ]
-    for activity_id in tqdm(
-        activity_ids_without_fingerprint, desc="Compute activity fingerprints"
-    ):
-        ts = repository.get_time_series(activity_id)
-        ts_hash = _compute_image_hash(ts)
-        fingerprints[activity_id] = ts_hash
-
-    distances = {}
-    if distances_path.exists():
-        with open(distances_path, "rb") as f:
-            distances = pickle.load(f)
-
-    for this in tqdm(
-        activity_ids_without_fingerprint, desc="Compute activity distances"
-    ):
-        for other in activity_ids:
-            distance = _hamming_distance(fingerprints[this], fingerprints[other])
-            add_distance(distances, this, other, distance)
-            add_distance(distances, other, this, distance)
-
-    with open(fingerprint_path, "wb") as f:
-        pickle.dump(fingerprints, f)
-    with open(distances_path, "wb") as f:
-        pickle.dump(distances, f)
+    with stored_object(fingerprint_path, {}) as fingerprints, stored_object(
+        distances_path, {}
+    ) as distances:
+        activity_ids = repository.activity_ids
+
+        activity_ids_without_fingerprint = [
+            activity_id
+            for activity_id in activity_ids
+            if activity_id not in fingerprints
+        ]
+        for activity_id in tqdm(
+            activity_ids_without_fingerprint, desc="Compute activity fingerprints"
+        ):
+            ts = repository.get_time_series(activity_id)
+            ts_hash = _compute_image_hash(ts)
+            fingerprints[activity_id] = ts_hash
+
+        for this in tqdm(
+            activity_ids_without_fingerprint, desc="Compute activity distances"
+        ):
+            for other in activity_ids:
+                distance = _hamming_distance(fingerprints[this], fingerprints[other])
+                add_distance(distances, this, other, distance)
+                add_distance(distances, other, this, distance)
 
 
 def asymmetric_activity_overlap(

geo_activity_playground/core/tasks.py

@@ -4,11 +4,32 @@
 import pickle
 from collections.abc import Iterable
 from typing import Any
+from typing import Generic
 from typing import Sequence
+from typing import TypeVar
 
 from geo_activity_playground.core.paths import cache_dir
 
 
+T = TypeVar("T")
+
+
+@contextlib.contextmanager
+def stored_object(path: pathlib.Path, default):
+    if path.exists():
+        with open(path, "rb") as f:
+            payload = pickle.load(f)
+    else:
+        payload = default
+
+    yield payload
+
+    temp_location = path.with_suffix(".tmp")
+    with open(temp_location, "wb") as f:
+        pickle.dump(payload, f)
+    temp_location.rename(path)
+
+
 def work_tracker_path(name: str) -> pathlib.Path:
     return cache_dir() / f"work-tracker-{name}.pickle"