Context manager for atomic file writes.
Ensures that files are either written completely or not at all,
preventing partial/corrupted files from power failures or crashes.
Usage
with AtomicPath(target_path, allow_override=False) as temp_path:
# Write to temp_path
with open(temp_path, 'w') as f:
f.write(data)
File is atomically moved to target_path on successful exit
Source code in inference/core/utils/file_system.py
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65 | class AtomicPath:
"""Context manager for atomic file writes.
Ensures that files are either written completely or not at all,
preventing partial/corrupted files from power failures or crashes.
Usage:
with AtomicPath(target_path, allow_override=False) as temp_path:
# Write to temp_path
with open(temp_path, 'w') as f:
f.write(data)
# File is atomically moved to target_path on successful exit
"""
def __init__(self, target_path: str, allow_override: bool = False):
self.target_path = target_path
self.allow_override = allow_override
self.temp_path: Optional[str] = None
self.temp_file = None
def __enter__(self) -> str:
ensure_write_is_allowed(
path=self.target_path, allow_override=self.allow_override
)
ensure_parent_dir_exists(path=self.target_path)
dir_name = os.path.dirname(os.path.abspath(self.target_path))
base_name = os.path.basename(self.target_path)
self.temp_file = tempfile.NamedTemporaryFile(
dir=dir_name, prefix=".tmp_", suffix="_" + base_name, delete=False
)
self.temp_path = self.temp_file.name
self.temp_file.close()
return self.temp_path
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is None:
try:
if os.name == "nt": # Windows
if os.path.exists(self.target_path):
os.remove(self.target_path)
os.rename(self.temp_path, self.target_path)
else: # POSIX
os.replace(self.temp_path, self.target_path)
except Exception:
try:
os.unlink(self.temp_path)
except OSError:
pass
raise
else:
# Error occurred - clean up temp file
try:
os.unlink(self.temp_path)
except OSError:
pass
return False # Don't suppress exceptions
|