Store the class metadata of fields in the class, to improve preformance

Cached data include,
- lookup table between groups and fields of "oneof" fields
- default value creator of each field
- type hint of each field

Author: jameslan

betterproto/__init__.py

@@ -120,7 +120,11 @@
 
 
 # Protobuf datetimes start at the Unix Epoch in 1970 in UTC.
-DATETIME_ZERO = datetime(1970, 1, 1, tzinfo=timezone.utc)
+def datetime_default_gen():
+    return datetime(1970, 1, 1, tzinfo=timezone.utc)
+
+
+DATETIME_ZERO = datetime_default_gen()
 
 
 class Casing(enum.Enum):
@@ -428,6 +432,57 @@ def parse_fields(value: bytes) -> Generator[ParsedField, None, None]:
 T = TypeVar("T", bound="Message")
 
 
+class ProtoClassMetadata:
+    cls: "Message"
+
+    def __init__(self, cls: "Message"):
+        self.cls = cls
+        by_field = {}
+        by_group = {}
+
+        for field in dataclasses.fields(cls):
+            meta = FieldMetadata.get(field)
+
+            if meta.group:
+                # This is part of a one-of group.
+                by_field[field.name] = meta.group
+
+                by_group.setdefault(meta.group, set()).add(field)
+
+        self.oneof_group_by_field = by_field
+        self.oneof_field_by_group = by_group
+
+    def __getattr__(self, item):
+        # Lazy init because forward reference classes may not be available at the beginning.
+        if item == 'default_gen':
+            defaults = {}
+            for field in dataclasses.fields(self.cls):
+                meta = FieldMetadata.get(field)
+                defaults[field.name] = self.cls._get_field_default_gen(field, meta)
+
+            self.default_gen = defaults  # __getattr__ won't be called next time
+            return defaults
+
+        if item == 'cls_by_field':
+            field_cls = {}
+            for field in dataclasses.fields(self.cls):
+                meta = FieldMetadata.get(field)
+                field_cls[field.name] = self.cls._type_hint(field.name)
+
+            self.cls_by_field = field_cls  # __getattr__ won't be called next time
+            return field_cls
+
+
+def make_protoclass(cls):
+    setattr(cls, "_betterproto", ProtoClassMetadata(cls))
+
+
+def protoclass(*args, **kwargs):
+    cls = dataclasses.dataclass(*args, **kwargs)
+    make_protoclass(cls)
+    return cls
+
+
 class Message(ABC):
     """
     A protobuf message base class. Generated code will inherit from this and
@@ -445,25 +500,20 @@ def __post_init__(self) -> None:
 
         # Set a default value for each field in the class after `__init__` has
         # already been run.
-        group_map: Dict[str, dict] = {"fields": {}, "groups": {}}
+        group_map: Dict[str, dataclasses.Field] = {}
         for field in dataclasses.fields(self):
             meta = FieldMetadata.get(field)
 
             if meta.group:
-                # This is part of a one-of group.
-                group_map["fields"][field.name] = meta.group
-
-                if meta.group not in group_map["groups"]:
-                    group_map["groups"][meta.group] = {"current": None, "fields": set()}
-                group_map["groups"][meta.group]["fields"].add(field)
+                group_map.setdefault(meta.group)
 
             if getattr(self, field.name) != PLACEHOLDER:
                 # Skip anything not set to the sentinel value
                 all_sentinel = False
 
                 if meta.group:
                     # This was set, so make it the selected value of the one-of.
-                    group_map["groups"][meta.group]["current"] = field
+                    group_map[meta.group] = field
 
                 continue
 
@@ -479,16 +529,17 @@ def __setattr__(self, attr: str, value: Any) -> None:
             # Track when a field has been set.
             self.__dict__["_serialized_on_wire"] = True
 
-        if attr in getattr(self, "_group_map", {}).get("fields", {}):
-            group = self._group_map["fields"][attr]
-            for field in self._group_map["groups"][group]["fields"]:
-                if field.name == attr:
-                    self._group_map["groups"][group]["current"] = field
-                else:
-                    super().__setattr__(
-                        field.name,
-                        self._get_field_default(field, FieldMetadata.get(field)),
-                    )
+        if hasattr(self, "_group_map"):  # __post_init__ had already run
+            if attr in self._betterproto.oneof_group_by_field:
+                group = self._betterproto.oneof_group_by_field[attr]
+                for field in self._betterproto.oneof_field_by_group[group]:
+                    if field.name == attr:
+                        self._group_map[group] = field
+                    else:
+                        super().__setattr__(
+                            field.name,
+                            self._get_field_default(field, FieldMetadata.get(field)),
+                        )
 
         super().__setattr__(attr, value)
 
@@ -510,7 +561,7 @@ def __bytes__(self) -> bytes:
             # currently set in a `oneof` group, so it must be serialized even
             # if the value is the default zero value.
             selected_in_group = False
-            if meta.group and self._group_map["groups"][meta.group]["current"] == field:
+            if meta.group and self._group_map[meta.group] == field:
                 selected_in_group = True
 
             serialize_empty = False
@@ -562,47 +613,49 @@ def __bytes__(self) -> bytes:
     # For compatibility with other libraries
     SerializeToString = __bytes__
 
-    def _type_hint(self, field_name: str) -> Type:
-        module = inspect.getmodule(self.__class__)
-        type_hints = get_type_hints(self.__class__, vars(module))
+    @classmethod
+    def _type_hint(cls, field_name: str) -> Type:
+        module = inspect.getmodule(cls)
+        type_hints = get_type_hints(cls, vars(module))
         return type_hints[field_name]
 
     def _cls_for(self, field: dataclasses.Field, index: int = 0) -> Type:
         """Get the message class for a field from the type hints."""
-        cls = self._type_hint(field.name)
+        cls = self._betterproto.cls_by_field[field.name]
         if hasattr(cls, "__args__") and index >= 0:
             cls = cls.__args__[index]
         return cls
 
     def _get_field_default(self, field: dataclasses.Field, meta: FieldMetadata) -> Any:
-        t = self._type_hint(field.name)
+        return self._betterproto.default_gen[field.name]()
+
+    @classmethod
+    def _get_field_default_gen(cls, field: dataclasses.Field, meta: FieldMetadata) -> Any:
+        t = cls._type_hint(field.name)
 
-        value: Any = 0
         if hasattr(t, "__origin__"):
             if t.__origin__ in (dict, Dict):
                 # This is some kind of map (dict in Python).
-                value = {}
+                return dict
             elif t.__origin__ in (list, List):
                 # This is some kind of list (repeated) field.
-                value = []
+                return list
             elif t.__origin__ == Union and t.__args__[1] == type(None):
                 # This is an optional (wrapped) field. For setting the default we
                 # really don't care what kind of field it is.
-                value = None
+                return type(None)
             else:
-                value = t()
+                return t
         elif issubclass(t, Enum):
             # Enums always default to zero.
-            value = 0
+            return int
         elif t == datetime:
             # Offsets are relative to 1970-01-01T00:00:00Z
-            value = DATETIME_ZERO
+            return datetime_default_gen
         else:
             # This is either a primitive scalar or another message type. Calling
             # it should result in its zero value.
-            value = t()
-
-        return value
+            return t
 
     def _postprocess_single(
         self, wire_type: int, meta: FieldMetadata, field: dataclasses.Field, value: Any
@@ -654,6 +707,7 @@ def _postprocess_single(
                     ],
                     bases=(Message,),
                 )
+                make_protoclass(Entry)
                 value = Entry().parse(value)
 
         return value
@@ -861,13 +915,13 @@ def serialized_on_wire(message: Message) -> bool:
 
 def which_one_of(message: Message, group_name: str) -> Tuple[str, Any]:
     """Return the name and value of a message's one-of field group."""
-    field = message._group_map["groups"].get(group_name, {}).get("current")
+    field = message._group_map.get(group_name)
     if not field:
         return ("", None)
     return (field.name, getattr(message, field.name))
 
 
-@dataclasses.dataclass
+@protoclass
 class _Duration(Message):
     # Signed seconds of the span of time. Must be from -315,576,000,000 to
     # +315,576,000,000 inclusive. Note: these bounds are computed from: 60
@@ -892,7 +946,7 @@ def delta_to_json(delta: timedelta) -> str:
         return ".".join(parts) + "s"
 
 
-@dataclasses.dataclass
+@protoclass
 class _Timestamp(Message):
     # Represents seconds of UTC time since Unix epoch 1970-01-01T00:00:00Z. Must
     # be from 0001-01-01T00:00:00Z to 9999-12-31T23:59:59Z inclusive.
@@ -942,47 +996,47 @@ def from_dict(self: T, value: Any) -> T:
         return self
 
 
-@dataclasses.dataclass
+@protoclass
 class _BoolValue(_WrappedMessage):
     value: bool = bool_field(1)
 
 
-@dataclasses.dataclass
+@protoclass
 class _Int32Value(_WrappedMessage):
     value: int = int32_field(1)
 
 
-@dataclasses.dataclass
+@protoclass
 class _UInt32Value(_WrappedMessage):
     value: int = uint32_field(1)
 
 
-@dataclasses.dataclass
+@protoclass
 class _Int64Value(_WrappedMessage):
     value: int = int64_field(1)
 
 
-@dataclasses.dataclass
+@protoclass
 class _UInt64Value(_WrappedMessage):
     value: int = uint64_field(1)
 
 
-@dataclasses.dataclass
+@protoclass
 class _FloatValue(_WrappedMessage):
     value: float = float_field(1)
 
 
-@dataclasses.dataclass
+@protoclass
 class _DoubleValue(_WrappedMessage):
     value: float = double_field(1)
 
 
-@dataclasses.dataclass
+@protoclass
 class _StringValue(_WrappedMessage):
     value: str = string_field(1)
 
 
-@dataclasses.dataclass
+@protoclass
 class _BytesValue(_WrappedMessage):
     value: bytes = bytes_field(1)