cleanup

.flake8

@@ -10,6 +10,7 @@ exclude =
   ./subtrees/
   ./venv/
   ./node_modules/
+  ./pyre
 application_import_names = backend
 import-order-style = edited
 

src/backend/common/models/cached_model.py

@@ -56,21 +56,22 @@ def __init__(self, *args, **kwargs):
         # constructors, so make sure we have a common set of properties defined
         self._fix_up_properties()
 
+    """
+    From the legacy NDB model implementation:
+    ```
+    def __getstate__(self):
+        return self._to_pb().Encode()
+
+    def __setstate__(self, serialized_pb):
+        pb = entity_pb.EntityProto(serialized_pb)
+        self.__init__()
+        self.__class__._from_pb(pb, set_key=False, ent=self)
+    ```
+
+    See https://github.com/googleapis/python-ndb/issues/587 about fixing upstream
+    """
+
     def __setstate__(self, state: Any) -> None:
-        """
-        From the legacy NDB model implementation:
-        ```
-        def __getstate__(self):
-            return self._to_pb().Encode()
-
-        def __setstate__(self, serialized_pb):
-            pb = entity_pb.EntityProto(serialized_pb)
-            self.__init__()
-            self.__class__._from_pb(pb, set_key=False, ent=self)
-        ```
-
-        See https://github.com/googleapis/python-ndb/issues/587 about fixing upstream
-        """
         pb = EntityProto()
         pb.MergePartialFromString(state)
 

src/backend/common/models/cached_query_result.py

@@ -1,9 +1,14 @@
 import io
 import pickle
+from typing import Any
 
 from google.cloud import ndb
 
 
+class ImportFixingPickler(pickle.Pickler):
+    pass
+
+
 class ImportFixingUnpickler(pickle.Unpickler):
     def find_class(self, module, name):
         renamed_module = module
@@ -14,16 +19,19 @@ def find_class(self, module, name):
 
 
 class ImportFixingPickleProperty(ndb.BlobProperty):
-    def _to_base_type(self, value):
+    def _to_base_type(self, value: Any) -> bytes:
         """Convert a value to the "base" value type for this property.
         Args:
             value (Any): The value to be converted.
         Returns:
             bytes: The pickled ``value``.
         """
-        return pickle.dumps(value, pickle.HIGHEST_PROTOCOL)
 
-    def _from_base_type(self, value):
+        file_obj = io.BytesIO()
+        ImportFixingPickler(file_obj, protocol=2, fix_imports=True).dump(value)
+        return file_obj.getvalue()
+
+    def _from_base_type(self, value: bytes) -> Any:
         """Convert a value from the "base" value type for this property.
         Args:
             value (bytes): The value to be converted.