Improve String Handling (#1132)

[etc/vim.git] / blib2to3 / pgen2 / grammar.py

diff --git a/blib2to3/pgen2/grammar.py b/blib2to3/pgen2/grammar.py
index c00cb2248d5676092c82d9da3f67d9b2f54f67ad..2882cdac89b2ac138151edd8d1791ea55db799f0 100644 (file)
--- a/blib2to3/pgen2/grammar.py
+++ b/blib2to3/pgen2/grammar.py
@@ -13,12 +13,20 @@ fallback token code OP, but the parser needs the actual token code.
 """
 
 # Python imports
-import collections
+import os
 import pickle
+import tempfile
+from typing import Any, Dict, List, Optional, Text, Tuple, TypeVar, Union
 
 # Local imports
 from . import token
 
+_P = TypeVar("_P", bound="Grammar")
+Label = Tuple[int, Optional[Text]]
+DFA = List[List[Tuple[int, int]]]
+DFAS = Tuple[DFA, Dict[int, int]]
+Path = Union[str, "os.PathLike[str]"]
+
 
 class Grammar(object):
     """Pgen parsing tables conversion class.
@@ -74,48 +82,74 @@ class Grammar(object):
 
     """
 
-    def __init__(self):
-        self.symbol2number = {}
-        self.number2symbol = {}
-        self.states = []
-        self.dfas = {}
-        self.labels = [(0, "EMPTY")]
-        self.keywords = {}
-        self.tokens = {}
-        self.symbol2label = {}
+    def __init__(self) -> None:
+        self.symbol2number: Dict[str, int] = {}
+        self.number2symbol: Dict[int, str] = {}
+        self.states: List[DFA] = []
+        self.dfas: Dict[int, DFAS] = {}
+        self.labels: List[Label] = [(0, "EMPTY")]
+        self.keywords: Dict[str, int] = {}
+        self.tokens: Dict[int, int] = {}
+        self.symbol2label: Dict[str, int] = {}
         self.start = 256
+        # Python 3.7+ parses async as a keyword, not an identifier
+        self.async_keywords = False
 
-    def dump(self, filename):
+    def dump(self, filename: Path) -> None:
         """Dump the grammar tables to a pickle file."""
-        with open(filename, "wb") as f:
-            pickle.dump(self.__dict__, f, pickle.HIGHEST_PROTOCOL)
 
-    def load(self, filename):
+        # mypyc generates objects that don't have a __dict__, but they
+        # do have __getstate__ methods that will return an equivalent
+        # dictionary
+        if hasattr(self, "__dict__"):
+            d = self.__dict__
+        else:
+            d = self.__getstate__()  # type: ignore
+
+        with tempfile.NamedTemporaryFile(
+            dir=os.path.dirname(filename), delete=False
+        ) as f:
+            pickle.dump(d, f, pickle.HIGHEST_PROTOCOL)
+        os.replace(f.name, filename)
+
+    def _update(self, attrs: Dict[str, Any]) -> None:
+        for k, v in attrs.items():
+            setattr(self, k, v)
+
+    def load(self, filename: Path) -> None:
         """Load the grammar tables from a pickle file."""
         with open(filename, "rb") as f:
             d = pickle.load(f)
-        self.__dict__.update(d)
+        self._update(d)
 
-    def loads(self, pkl):
+    def loads(self, pkl: bytes) -> None:
         """Load the grammar tables from a pickle bytes object."""
-        self.__dict__.update(pickle.loads(pkl))
+        self._update(pickle.loads(pkl))
 
-    def copy(self):
+    def copy(self: _P) -> _P:
         """
         Copy the grammar.
         """
         new = self.__class__()
-        for dict_attr in ("symbol2number", "number2symbol", "dfas", "keywords",
-                          "tokens", "symbol2label"):
+        for dict_attr in (
+            "symbol2number",
+            "number2symbol",
+            "dfas",
+            "keywords",
+            "tokens",
+            "symbol2label",
+        ):
             setattr(new, dict_attr, getattr(self, dict_attr).copy())
         new.labels = self.labels[:]
         new.states = self.states[:]
         new.start = self.start
+        new.async_keywords = self.async_keywords
         return new
 
-    def report(self):
+    def report(self) -> None:
         """Dump the grammar tables to standard output, for debugging."""
         from pprint import pprint
+
         print("s2n")
         pprint(self.symbol2number)
         print("n2s")
@@ -179,6 +213,7 @@ opmap_raw = """
 // DOUBLESLASH
 //= DOUBLESLASHEQUAL
 -> RARROW
+:= COLONEQUAL
 """
 
 opmap = {}