From: Anthony Sottile <asottile@umich.edu>
Date: Sat, 19 Jan 2019 04:59:17 +0000 (-0800)
Subject: Atomically write cache files (#674)
X-Git-Url: https://git.madduck.net/etc/vim.git/commitdiff_plain/4d3107233f18452d89e3caf148a098849a16b364?ds=inline;pf=etc

Atomically write cache files (#674)
---

diff --git a/black.py b/black.py
index 57d5c83..3269322 100644
--- a/black.py
+++ b/black.py
@@ -14,6 +14,7 @@ import pickle
 import re
 import signal
 import sys
+import tempfile
 import tokenize
 from typing import (
     Any,
@@ -3647,11 +3648,11 @@ def write_cache(
     """Update the cache file."""
     cache_file = get_cache_file(line_length, mode)
     try:
-        if not CACHE_DIR.exists():
-            CACHE_DIR.mkdir(parents=True)
+        CACHE_DIR.mkdir(parents=True, exist_ok=True)
         new_cache = {**cache, **{src.resolve(): get_cache_info(src) for src in sources}}
-        with cache_file.open("wb") as fobj:
-            pickle.dump(new_cache, fobj, protocol=pickle.HIGHEST_PROTOCOL)
+        with tempfile.NamedTemporaryFile(dir=str(cache_file.parent), delete=False) as f:
+            pickle.dump(new_cache, f, protocol=pickle.HIGHEST_PROTOCOL)
+        os.replace(f.name, cache_file)
     except OSError:
         pass
 
diff --git a/blib2to3/pgen2/grammar.py b/blib2to3/pgen2/grammar.py
index a1da546..b20b923 100644
--- a/blib2to3/pgen2/grammar.py
+++ b/blib2to3/pgen2/grammar.py
@@ -13,7 +13,9 @@ fallback token code OP, but the parser needs the actual token code.
 """
 
 # Python imports
+import os.path
 import pickle
+import tempfile
 
 # Local imports
 from . import token
@@ -86,8 +88,9 @@ class Grammar(object):
 
     def dump(self, filename):
         """Dump the grammar tables to a pickle file."""
-        with open(filename, "wb") as f:
+        with tempfile.NamedTemporaryFile(dir=os.path.dirname(filename), delete=False) as f:
             pickle.dump(self.__dict__, f, pickle.HIGHEST_PROTOCOL)
+        os.replace(f.name, filename)
 
     def load(self, filename):
         """Load the grammar tables from a pickle file."""