ylib2to3/pgen2/grammar: Fix Grammar.dump for parallel use with Gramma…

…r.load (#1243)

Previously, bad timing could make another process run into reading a
half-written pickle cache file, and thus fail like this:

> Traceback (most recent call last):
>   File "[..]/bin/yapf", line 5, in <module>
>     from yapf import run_main
>   File "[..]/lib/python3.11/site-packages/yapf/__init__.py", line 41, in <module>
>     from yapf.yapflib import yapf_api
>   File "[..]/lib/python3.11/site-packages/yapf/yapflib/yapf_api.py", line 38, in <module>
>     from yapf.pyparser import pyparser
>   File "[..]/lib/python3.11/site-packages/yapf/pyparser/pyparser.py", line 44, in <module>
>     from yapf.yapflib import format_token
>   File "[..]/lib/python3.11/site-packages/yapf/yapflib/format_token.py", line 23, in <module>
>     from yapf.pytree import pytree_utils
>   File "[..]/lib/python3.11/site-packages/yapf/pytree/pytree_utils.py", line 30, in <module>
>     from yapf_third_party._ylib2to3 import pygram
>   File "[..]/lib/python3.11/site-packages/yapf_third_party/_ylib2to3/pygram.py", line 29, in <module>
>     python_grammar = driver.load_grammar(_GRAMMAR_FILE)
>                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
>   File "[..]/lib/python3.11/site-packages/yapf_third_party/_ylib2to3/pgen2/driver.py", line 252, in load_grammar
>     g.load(gp)
>   File "[..]/lib/python3.11/site-packages/yapf_third_party/_ylib2to3/pgen2/grammar.py", line 95, in load
>     d = pickle.load(f)
>         ^^^^^^^^^^^^^^
> EOFError: Ran out of input

third_party/yapf_third_party/_ylib2to3/pgen2/grammar.py

@@ -12,7 +12,9 @@
 """
 
 # Python imports
+import os
 import pickle
+import tempfile
 
 # Local imports
 from . import token
@@ -86,8 +88,39 @@ def __init__(self):
 
   def dump(self, filename):
     """Dump the grammar tables to a pickle file."""
-    with open(filename, 'wb') as f:
-      pickle.dump(self.__dict__, f, pickle.HIGHEST_PROTOCOL)
+    # NOTE:
+    # - We're writing a tempfile first so that there is no chance
+    #   for someone to read a half-written file from this very spot
+    #   while we're were not done writing.
+    # - We're using ``os.rename`` to sure not copy data around (which
+    #   would get us back to square one with a reading-half-written file
+    #   race condition).
+    # - We're making the tempfile go to the same directory as the eventual
+    #   target ``filename`` so that there is no chance of failing from
+    #   cross-file-system renames in ``os.rename``.
+    # - We're using the same prefix and suffix for the tempfile so if we
+    #   ever have to leave a tempfile around for failure of deletion,
+    #   it will have a reasonable filename extension and its name will help
+    #   explain is nature.
+    tempfile_dir = os.path.dirname(filename)
+    tempfile_prefix, tempfile_suffix = os.path.splitext(filename)
+    with tempfile.NamedTemporaryFile(
+        mode='wb',
+        suffix=tempfile_suffix,
+        prefix=tempfile_prefix,
+        dir=tempfile_dir,
+        delete=False) as f:
+      pickle.dump(self.__dict__, f.file, pickle.HIGHEST_PROTOCOL)
+      try:
+        os.rename(f.name, filename)
+      except OSError:
+        # This makes sure that we do not leave the tempfile around
+        # unless we have to...
+        try:
+          os.remove(f.name)
+        except OSError:
+          pass
+        raise
 
   def load(self, filename):
     """Load the grammar tables from a pickle file."""