start of pb performance refactor

Author: craiglabenz

proto/fields.py

@@ -114,6 +114,21 @@ def descriptor(self):
         # Return the descriptor.
         return self._descriptor
 
+    def contribute_to_class(self, cls, name: str):
+
+        class _Descriptor:
+            def __init__(self, name: str):
+                self.name = f'_{name}'
+
+            def __set__(self, instance, value):
+                setattr(instance, self.name, value)
+                instance._mark_dirty()
+
+            def __get__(self, instance, owner):
+                return getattr(instance, self.name, None)
+
+        setattr(cls, name, _Descriptor(name))
+
     @property
     def name(self) -> str:
         """Return the name of the field."""

proto/message.py

@@ -15,6 +15,7 @@
 import collections
 import collections.abc
 import copy
+import inspect
 import re
 from typing import List, Type
 
@@ -31,6 +32,11 @@
 from proto.primitives import ProtoType
 
 
+def _has_contribute_to_class(value):
+    # Only call contribute_to_class() if it's bound.
+    return not inspect.isclass(value) and hasattr(value, 'contribute_to_class')
+
+
 class MessageMeta(type):
     """A metaclass for building and registering Message subclasses."""
 
@@ -105,6 +111,7 @@ def __new__(mcls, name, bases, attrs):
         # Okay, now we deal with all the rest of the fields.
         # Iterate over all the attributes and separate the fields into
         # their own sequence.
+        contributable_attrs = {}
         fields = []
         new_attrs = {}
         oneofs = collections.OrderedDict()
@@ -127,6 +134,9 @@ def __new__(mcls, name, bases, attrs):
                 "package": package,
             }
 
+            if _has_contribute_to_class(field):
+                contributable_attrs[key] = field
+
             # Add the field to the list of fields.
             fields.append(field)
             # If this field is part of a "oneof", ensure the oneof itself
@@ -248,6 +258,9 @@ def __new__(mcls, name, bases, attrs):
         # Run the superclass constructor.
         cls = super().__new__(mcls, name, bases, new_attrs)
 
+        for field_name, field in contributable_attrs.items():
+            cls.add_to_class(field_name, field)
+
         # The info class and fields need a reference to the class just created.
         cls._meta.parent = cls
         for field in cls._meta.fields.values():
@@ -269,6 +282,12 @@ def __new__(mcls, name, bases, attrs):
     def __prepare__(mcls, name, bases, **kwargs):
         return collections.OrderedDict()
 
+    def add_to_class(cls, name, value):
+        if _has_contribute_to_class(value):
+            value.contribute_to_class(cls, name)
+        else:
+            setattr(cls, name, value)
+
     @property
     def meta(cls):
         return cls._meta
@@ -313,6 +332,9 @@ def serialize(cls, instance) -> bytes:
         Returns:
             bytes: The serialized representation of the protocol buffer.
         """
+        if instance.is_dirty:
+            new_pb_values = instance._map_from_fields()
+            instance._update_pb(new_pb_values)
         return cls.pb(instance, coerce=True).SerializeToString()
 
     def deserialize(cls, payload: bytes) -> "Message":
@@ -445,6 +467,13 @@ class Message(metaclass=MessageMeta):
             message.
     """
 
+    CLEAN_ATTRS = '_clean_attrs'
+
+    @property
+    def is_dirty(self) -> bool:
+        """Default state is "clean", so the attr not being set indicates no problems."""
+        return not getattr(self, self.CLEAN_ATTRS, True)
+
     def __init__(self, mapping=None, *, ignore_unknown_fields=False, **kwargs):
         # We accept several things for `mapping`:
         #   * An instance of this class.
@@ -454,7 +483,7 @@ def __init__(self, mapping=None, *, ignore_unknown_fields=False, **kwargs):
         if mapping is None:
             if not kwargs:
                 # Special fast path for empty construction.
-                super().__setattr__("_pb", self._meta.pb())
+                # super().__setattr__("_pb", self._meta.pb())
                 return
 
             mapping = kwargs
@@ -488,10 +517,10 @@ def __init__(self, mapping=None, *, ignore_unknown_fields=False, **kwargs):
                 % (self.__class__.__name__, mapping,)
             )
 
-        params = {}
+        # params = {}
         # Update the mapping to address any values that need to be
         # coerced.
-        marshal = self._meta.marshal
+        # marshal = self._meta.marshal
         for key, value in mapping.items():
             try:
                 pb_type = self._meta.fields[key].pb_type
@@ -503,12 +532,14 @@ def __init__(self, mapping=None, *, ignore_unknown_fields=False, **kwargs):
                     "Unknown field for {}: {}".format(self.__class__.__name__, key)
                 )
 
-            pb_value = marshal.to_proto(pb_type, value)
-            if pb_value is not None:
-                params[key] = pb_value
+            # pb_value = marshal.to_proto(pb_type, value)
+            # if pb_value is not None:
+            #     params[key] = pb_value
+            setattr(self, key, value)
 
         # Create the internal protocol buffer.
-        super().__setattr__("_pb", self._meta.pb(**params))
+        super().__setattr__("_pb", self._meta.pb())
+        self._mark_dirty()
 
     def __bool__(self):
         """Return True if any field is truthy, False otherwise."""
@@ -609,28 +640,72 @@ def __ne__(self, other):
         return not self == other
 
     def __repr__(self):
-        return repr(self._pb)
+        if (getattr(self, '_pb'), None) is not None:
+            return repr(self._pb)
+        return repr(self)
+
+    # def __setattr__(self, key, value):
+    #     """Set the value on the given field.
+
+    #     For well-known protocol buffer types which are marshalled, either
+    #     the protocol buffer object or the Python equivalent is accepted.
+    #     """
+    #     if key[0] == "_" or key in self.__dict__:
+    #         return super().__setattr__(key, value)
+    #     marshal = self._meta.marshal
+    #     pb_type = self._meta.fields[key].pb_type
+    #     pb_value = marshal.to_proto(pb_type, value)
+
+    #     # Clear the existing field.
+    #     # This is the only way to successfully write nested falsy values,
+    #     # because otherwise MergeFrom will no-op on them.
+    #     self._pb.ClearField(key)
+
+    #     # Merge in the value being set.
+    #     if pb_value is not None:
+    #         self._pb.MergeFrom(self._meta.pb(**{key: pb_value}))
 
-    def __setattr__(self, key, value):
-        """Set the value on the given field.
-
-        For well-known protocol buffer types which are marshalled, either
-        the protocol buffer object or the Python equivalent is accepted.
+    @property
+    def _pb(self):
+        if not hasattr(self, '_cached_pb'):
+            self._cached_pb = self._meta.pb()
+        return self._cached_pb
+
+    @_pb.setter
+    def _pb(self, value):
+        self._cached_pb = value
+
+    def _map_from_fields(self) -> dict:
+        _map = {}
+        for field_name in self._meta.fields.keys():
+            _map[field_name] = getattr(self, field_name)
+        return _map
+
+    def _update_pb(self, values: dict):
+        """Batch update of inner _pb, used before serialization
         """
-        if key[0] == "_":
-            return super().__setattr__(key, value)
-        marshal = self._meta.marshal
-        pb_type = self._meta.fields[key].pb_type
-        pb_value = marshal.to_proto(pb_type, value)
-
-        # Clear the existing field.
-        # This is the only way to successfully write nested falsy values,
-        # because otherwise MergeFrom will no-op on them.
-        self._pb.ClearField(key)
+        _marshalled = {}
+        for key, value in values.items():
+            marshal = self._meta.marshal
+            pb_type = self._meta.fields[key].pb_type
+            pb_value = marshal.to_proto(pb_type, value)
+            _marshalled[key] = pb_value
+
+            # Clear the existing field.
+            # This is the only way to successfully write nested falsy values,
+            # because otherwise MergeFrom will no-op on them.
+            self._pb.ClearField(key)
 
         # Merge in the value being set.
         if pb_value is not None:
-            self._pb.MergeFrom(self._meta.pb(**{key: pb_value}))
+            self._pb.MergeFrom(self._meta.pb(**_marshalled))
+        self._mark_clean()
+
+    def _mark_clean(self):
+        self._clean_attrs = True
+
+    def _mark_dirty(self):
+        self._clean_attrs = False
 
 
 class _MessageInfo: