Skip to content

Commit

Permalink
Delete local context directories on exit (#524)
Browse files Browse the repository at this point in the history
  • Loading branch information
tomwhite authored Jul 30, 2024
1 parent cb47ea1 commit debb1d8
Showing 1 changed file with 18 additions and 0 deletions.
18 changes: 18 additions & 0 deletions cubed/core/plan.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import atexit
import inspect
import shutil
import tempfile
import uuid
from datetime import datetime
Expand All @@ -19,6 +21,16 @@
# A unique ID with sensible ordering, used for making directory names
CONTEXT_ID = f"cubed-{datetime.now().strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4()}"

# Delete local context dirs when Python exits
CONTEXT_DIRS = set()


def delete_on_exit(context_dir: str) -> None:
if context_dir not in CONTEXT_DIRS and context_dir.startswith("/"):
atexit.register(lambda: shutil.rmtree(context_dir, ignore_errors=True))
CONTEXT_DIRS.add(context_dir)


sym_counter = 0


Expand Down Expand Up @@ -436,10 +448,16 @@ def arrays_to_plan(*arrays):


def new_temp_path(name, suffix=".zarr", spec=None):
"""Return a string path for a temporary file path, which may be local or remote.
Note that this function does not create the file or any directories (and they
may never be created, if for example the file doesn't need to be materialized).
"""
work_dir = spec.work_dir if spec is not None else None
if work_dir is None:
work_dir = tempfile.gettempdir()
context_dir = join_path(work_dir, CONTEXT_ID)
delete_on_exit(context_dir)
return join_path(context_dir, f"{name}{suffix}")


Expand Down

0 comments on commit debb1d8

Please sign in to comment.