responded to more code review

proto/fields.py

@@ -18,6 +18,7 @@
 from proto.datetime_helpers import DatetimeWithNanoseconds
 from proto.marshal.rules import wrappers
 from proto.marshal.rules.dates import DurationRule
+from proto.marshal.collections.maps import MapComposite
 from proto.utils import cached_property
 from typing import Any, Callable, Optional, Union
 
@@ -153,16 +154,11 @@ def pb_type(self):
         return self.message
 
     @property
-    def can_get_natively(self) -> bool:
-        if self.proto_type == ProtoType.MESSAGE and self.message == struct_pb2.Value:
-            return False
-        return True
-
-    @property
-    def can_set_natively(self) -> bool:
-        if self.proto_type == ProtoType.MESSAGE and self.message == struct_pb2.Value:
-            return False
-        return True
+    def can_represent_natively(self) -> bool:
+        return not (
+            self.proto_type == ProtoType.MESSAGE and
+            self.message == struct_pb2.Value
+        )
 
     def contribute_to_class(self, cls, name: str):
         """Attaches a descriptor to the top-level proto.Message class, so that attribute
@@ -288,6 +284,9 @@ def __init__(self, name: str, *, cls, set_coercion: Optional[Callable] = None):
 
         # simple types coercion for setting attributes
         # (e.g., bytes -> str if our type is string, but we are supplied bytes)
+        # the signature of `set_coercion` is dependent on the field's data types
+        # and is always handled by `contribute_to_class` which pairs data types
+        # to appropriate write-time coercions.
         self._set_coercion: Optional[Callable] = set_coercion
         self.cls = cls
 
@@ -349,7 +348,7 @@ class MyMessage(proto.Message):
         # instances receive attribute updates, immediately syncing those values to the underlying
         # pb2 instance is sufficient.
         always_commit: bool = getattr(instance, '_always_commit', False)
-        if always_commit or not self.field.can_set_natively:
+        if always_commit or not self.field.can_represent_natively:
             pb_value = instance._meta.marshal.to_proto(self.field.pb_type, value)
             _pb = instance._meta.pb(**{self.original_name: pb_value})
             instance._pb.ClearField(self.original_name)
@@ -379,26 +378,40 @@ class MyMessage(proto.Message):
             my_message = MyMessage(name="Frodo")
             print(my_message.name)
 
-        In the above scenario, `__get__` is called with "my_message" passed as `instance`.
+        In the above scenario, `__get__` is called with "my_message" passed as
+        `instance`.
         """
         # If `instance` is None, then we are accessing this field directly
         # off the class itself instead of off an instance.
         if instance is None:
             return self.original_name
 
         value = getattr(instance, self.instance_attr_name, _none)
-        if self.field.can_get_natively and value is not _none:
+        is_map: bool = isinstance(value, MapComposite)
+
+        # Return any values that do not require immediate rehydration.
+        # A few notes:
+        #   * primitives are simple, and so can be returned
+        #   * `Messages` are already Pythonic, and so can be returned
+        #   * `Values` are wrappers and so have to be unwrapped
+        #     * The exception to this is MapComposites, which have the same
+        #       types as Values, but which handle their own field caching and
+        #       thus can be returned when pulled off the `instance`.
+        if value is not _none and (is_map or self.field.can_represent_natively):
             return value
 
-        # For the most part, only primitive values can be returned natively, meaning
-        # this is either a Message itself, in which case, since we're dealing with the
-        # underlying pb object, we need to sync all deferred fields.
-        # This is functionally a no-op if no fields have been deferred.
+        # For the most part, only primitive values can be returned natively,
+        # meaning this is either a Message itself, in which case, since we're
+        # dealing with the underlying pb object, we need to sync all deferred
+        # fields. This is functionally a no-op if no fields have been deferred.
         if hasattr(value, '_update_pb'):
             value._update_pb()
 
         pb_value = getattr(instance._pb, self.original_name, None)
-        value = instance._meta.marshal.to_python(self.field.pb_type, pb_value, absent=self.original_name not in instance)
+        value = instance._meta.marshal.to_python(
+            self.field.pb_type, pb_value,
+            absent=self.original_name not in instance,
+        )
 
         setattr(instance, self.instance_attr_name, value)
         return value

proto/marshal/collections/maps.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import collections
+from typing import Dict, Set
 
 import proto
 from proto.utils import cached_property
@@ -41,26 +42,59 @@ def __init__(self, sequence, *, marshal):
         """
         self._pb = sequence
         self._marshal = marshal
+        self._item_cache: Dict = {}
+        self._stale_keys: Set[str] = set()
 
     def __contains__(self, key):
         # Protocol buffers is so permissive that querying for the existence
         # of a key will in of itself create it.
         #
         # By taking a tuple of the keys and querying that, we avoid sending
         # the lookup to protocol buffers and therefore avoid creating the key.
+        if key in self._item_cache:
+            return True
         return key in tuple(self.keys())
 
     def __getitem__(self, key):
         # We handle raising KeyError ourselves, because otherwise protocol
         # buffers will create the key if it does not exist.
-        if key not in self:
-            raise KeyError(key)
-        obj = self._marshal.to_python(self._pb_type, self.pb[key])
-        if isinstance(obj, proto.Message):
-            obj._always_commit = True
-        return obj
+        value = self._item_cache.get(key, _Empty.shared)
+
+        if isinstance(value, _Empty):
+            if key not in self:
+                raise KeyError(key)
+            value = self._marshal.to_python(self._pb_type, self.pb[key])
+
+            # This is the first domino in a hacky workaround that is completed
+            # in `fields._FieldDescriptor.__set__`. Because of the by-value nature
+            # of protobufs (which conflicts with the by-reference nature of Python),
+            # proto-plus objects that are yielded from MapFields must immediately
+            # write to their internal pb2 object whenever their fields are updated.
+            # This is a new requirement as always writing to a proto-plus object's
+            # inner pb2 protobuf used to be the default, but has been moved to a
+            # lazy-syncing system for performance reasons.
+            if isinstance(value, proto.Message):
+                value._always_commit = True
+
+            self._item_cache[key] = value
+
+        return value
 
     def __setitem__(self, key, value):
+        self._item_cache[key] = value
+        self._stale_keys.add(key)
+        # self._sync_key(key, value)
+
+    def _sync_all_keys(self):
+        for key in self._stale_keys:
+            self._sync_key(key)
+        self._stale_keys.clear()
+
+    def _sync_key(self, key, value = None):
+        value = value or self._item_cache.pop(key, None)
+        if value is None:
+            self.pb.pop(key)
+            return
         pb_value = self._marshal.to_proto(self._pb_type, value, strict=True)
 
         # Directly setting a key is not allowed; however, protocol buffers
@@ -73,14 +107,27 @@ def __setitem__(self, key, value):
         self.pb[key].MergeFrom(pb_value)
 
     def __delitem__(self, key):
-        self.pb.pop(key)
+        self._item_cache.pop(key, None)
+        self._stale_keys.add(key)
+        # self.pb.pop(key)
 
     def __len__(self):
-        return len(self.pb)
+        _all_keys = set(list(self._item_cache.keys()))
+        _all_keys = _all_keys.union(list(self.pb.keys()))
+        return len(_all_keys)
+        # return len(self.pb)
 
     def __iter__(self):
+        self._sync_all_keys()
         return iter(self.pb)
 
     @property
     def pb(self):
         return self._pb
+
+
+class _Empty:
+    pass
+
+
+_Empty.shared = _Empty()

proto/message.py

@@ -483,7 +483,6 @@ def __init__(self, mapping=None, *, ignore_unknown_fields=False, **kwargs):
         # underlying `_pb`
         self._stale_fields: Set[str] = set()
 
-
         # We accept several things for `mapping`:
         #   * An instance of this class.
         #   * An instance of the underlying protobuf descriptor class.
@@ -592,7 +591,6 @@ def __contains__(self, key):
         except ValueError:
             return bool(pb_value)
 
-
     def __eq__(self, other):
         """Return True if the messages are equal, False otherwise."""
         # If these are the same type, use internal protobuf's equality check.
@@ -633,7 +631,8 @@ def _mark_pb_stale(self, field_name: str):
         self._stale_fields.add(field_name)
 
     def _mark_pb_synced(self):
-        self._stale_fields = set()
+        if hasattr(self, '_stale_fields'):
+            self._stale_fields.clear()
 
     def _update_nested_pb(self):
         """When it is time to serialize a pb2 object, it does not do to sync just ourselves -
@@ -643,8 +642,11 @@ def _update_nested_pb(self):
         for field_name, field in self._meta.fields.items():
             if field.proto_type == proto.MESSAGE:
                 obj = getattr(self, field_name, None)
+                # Traversing the tree of objects eventually encounters primitives
+                # that do not have these functions
                 if obj and hasattr(obj, '_update_pb'):
                     obj._update_pb()
+                # Same here - many values will not have these methods
                 if obj and hasattr(obj, '_update_nested_pb'):
                     obj._update_nested_pb()
 

tests/test_fields_composite.py

@@ -24,6 +24,7 @@ class Spam(proto.Message):
         foo = proto.Field(proto.MESSAGE, number=1, message=Foo)
         eggs = proto.Field(proto.BOOL, number=2)
 
+
     spam = Spam(foo=Foo(bar="str", baz=42))
     assert spam.foo.bar == "str"
     assert spam.foo.baz == 42