Atomicized data.json.write()

Author: adamlui

remove-json-keys/src/remove_json_keys/lib/data/file.py

@@ -1,3 +1,19 @@
+def atomic_write(file_path, data, encoding='utf-8'): # to prevent TOCTOU
+    import os
+    from pathlib import Path
+
+    file_path = Path(file_path)
+    file_path.parent.mkdir(parents=True, exist_ok=True)
+    tmp_path = file_path.parent / f'.{file_path.name}.tmp'
+
+    try:
+        with open(tmp_path, 'w', encoding=encoding) as file:
+            file.write(data) ; file.flush() ; os.fsync(file.fileno())
+        os.replace(tmp_path, file_path) # atomic rename
+    except Exception:
+        if tmp_path.exists() : tmp_path.unlink()
+        raise
+
 def read(file_path, encoding='utf-8'):
     with open(file_path, 'r', encoding=encoding) as file:
         return file.read()

remove-json-keys/src/remove_json_keys/lib/data/json.py

@@ -43,19 +43,20 @@ def remove_keys(json_path, keys):
         files_processed_cnt += 1
     return keys_removed, keys_skipped, files_processed_cnt
 
-def write(file_path, data, encoding='utf-8', ensure_ascii=False, style='pretty'):
+def write(file_path, data, encoding='utf-8', ensure_ascii=False, style='pretty', atomic=True):
     Path(file_path).parent.mkdir(parents=True, exist_ok=True)
-    with open(file_path, 'w', encoding=encoding) as file:
-        if style == 'pretty': # single key/val spans multi-lines
-            json.dump(data, file, indent=2, ensure_ascii=ensure_ascii)
-        elif style == 'compact': # single key/val per line
-            file.write('{\n')
-            items = list(data.items())
-            for idx, (key, val) in enumerate(items):
-                line_end = ',' if idx < len(items) -1 else ''
-                inner = f'{{ {json.dumps(val, ensure_ascii=ensure_ascii)[1:-1]} }}'
-                file.write(f'  "{key}": {inner}{line_end}\n')
-            file.write('}')
-        else: # minified to single line
-            json.dump(data, file, separators=(',', ':'), ensure_ascii=ensure_ascii)
-        file.write('\n') # trailing newline
+    if style == 'pretty': # single key/val spans multi-lines
+        json_str = json.dumps(data, indent=2, ensure_ascii=ensure_ascii)
+    elif style == 'compact': # single key/val per line
+        lines = ['{']
+        items = list(data.items())
+        for idx, (key, val) in enumerate(items):
+            line_end = ',' if idx < len(items) -1 else ''
+            inner = f'{{ {json.dumps(val, ensure_ascii=ensure_ascii)[1:-1]} }}'
+            lines.append(f'  "{key}": {inner}{line_end}')
+        lines.append('}')
+        json_str = '\n'.join(lines)
+    else:  # minified
+        json_str = json.dumps(data, separators=(',', ':'), ensure_ascii=ensure_ascii)
+    json_str += '\n'
+    getattr(file, 'atomic_write' if atomic else 'write')(file_path, json_str, encoding=encoding)

translate-messages/src/translate_messages/lib/data/file.py

@@ -1,3 +1,19 @@
+def atomic_write(file_path, data, encoding='utf-8'): # to prevent TOCTOU
+    import os
+    from pathlib import Path
+
+    file_path = Path(file_path)
+    file_path.parent.mkdir(parents=True, exist_ok=True)
+    tmp_path = file_path.parent / f'.{file_path.name}.tmp'
+
+    try:
+        with open(tmp_path, 'w', encoding=encoding) as file:
+            file.write(data) ; file.flush() ; os.fsync(file.fileno())
+        os.replace(tmp_path, file_path) # atomic rename
+    except Exception:
+        if tmp_path.exists() : tmp_path.unlink()
+        raise
+
 def read(file_path, encoding='utf-8'):
     with open(file_path, 'r', encoding=encoding) as file:
         return file.read()

translate-messages/src/translate_messages/lib/data/json.py

@@ -23,19 +23,21 @@ def read(file_path, encoding='utf-8'):
     with open(file_path, 'r', encoding=encoding) as file:
         return json5.load(file)
 
-def write(file_path, data, encoding='utf-8', ensure_ascii=False, style='pretty'):
+def write(file_path, data, encoding='utf-8', ensure_ascii=False, style='pretty', atomic=True):
+    from . import file
     Path(file_path).parent.mkdir(parents=True, exist_ok=True)
-    with open(file_path, 'w', encoding=encoding) as file:
-        if style == 'pretty': # single key/val spans multi-lines
-            json.dump(data, file, indent=2, ensure_ascii=ensure_ascii)
-        elif style == 'compact': # single key/val per line
-            file.write('{\n')
-            items = list(data.items())
-            for idx, (key, val) in enumerate(items):
-                line_end = ',' if idx < len(items) -1 else ''
-                inner = f'{{ {json.dumps(val, ensure_ascii=ensure_ascii)[1:-1]} }}'
-                file.write(f'  "{key}": {inner}{line_end}\n')
-            file.write('}')
-        else: # minified to single line
-            json.dump(data, file, separators=(',', ':'), ensure_ascii=ensure_ascii)
-        file.write('\n') # trailing newline
+    if style == 'pretty': # single key/val spans multi-lines
+        json_str = json.dumps(data, indent=2, ensure_ascii=ensure_ascii)
+    elif style == 'compact': # single key/val per line
+        lines = ['{']
+        items = list(data.items())
+        for idx, (key, val) in enumerate(items):
+            line_end = ',' if idx < len(items) -1 else ''
+            inner = f'{{ {json.dumps(val, ensure_ascii=ensure_ascii)[1:-1]} }}'
+            lines.append(f'  "{key}": {inner}{line_end}')
+        lines.append('}')
+        json_str = '\n'.join(lines)
+    else:  # minified
+        json_str = json.dumps(data, separators=(',', ':'), ensure_ascii=ensure_ascii)
+    json_str += '\n'
+    getattr(file, 'atomic_write' if atomic else 'write')(file_path, json_str, encoding=encoding)