Use external package betterproto-rust-codec for better (de-)seriali…

…zation performance (#545) * optionally use betterproto-rust-codec * monkey patch `parse` and `__bytes__`
danielgtaylor · Dec 7, 2023 · d34b169 · d34b169
1 parent bd7de20
commit d34b169
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 33 deletions.
diff --git a/benchmarks/benchmarks.py b/benchmarks/benchmarks.py
@@ -6,32 +6,32 @@
 
 @dataclass
 class TestMessage(betterproto.Message):
-    foo: int = betterproto.uint32_field(0)
-    bar: str = betterproto.string_field(1)
-    baz: float = betterproto.float_field(2)
+    foo: int = betterproto.uint32_field(1)
+    bar: str = betterproto.string_field(2)
+    baz: float = betterproto.float_field(3)
 
 
 @dataclass
 class TestNestedChildMessage(betterproto.Message):
-    str_key: str = betterproto.string_field(0)
-    bytes_key: bytes = betterproto.bytes_field(1)
-    bool_key: bool = betterproto.bool_field(2)
-    float_key: float = betterproto.float_field(3)
-    int_key: int = betterproto.uint64_field(4)
+    str_key: str = betterproto.string_field(1)
+    bytes_key: bytes = betterproto.bytes_field(2)
+    bool_key: bool = betterproto.bool_field(3)
+    float_key: float = betterproto.float_field(4)
+    int_key: int = betterproto.uint64_field(5)
 
 
 @dataclass
 class TestNestedMessage(betterproto.Message):
-    foo: TestNestedChildMessage = betterproto.message_field(0)
-    bar: TestNestedChildMessage = betterproto.message_field(1)
-    baz: TestNestedChildMessage = betterproto.message_field(2)
+    foo: TestNestedChildMessage = betterproto.message_field(1)
+    bar: TestNestedChildMessage = betterproto.message_field(2)
+    baz: TestNestedChildMessage = betterproto.message_field(3)
 
 
 @dataclass
 class TestRepeatedMessage(betterproto.Message):
-    foo_repeat: List[str] = betterproto.string_field(0)
-    bar_repeat: List[int] = betterproto.int64_field(1)
-    baz_repeat: List[bool] = betterproto.bool_field(2)
+    foo_repeat: List[str] = betterproto.string_field(1)
+    bar_repeat: List[int] = betterproto.int64_field(2)
+    baz_repeat: List[bool] = betterproto.bool_field(3)
 
 
 class BenchMessage:
@@ -44,25 +44,14 @@ def setup(self):
         self.instance_filled_bytes = bytes(self.instance_filled)
         self.instance_filled_nested = TestNestedMessage(
             TestNestedChildMessage("foo", bytearray(b"test1"), True, 0.1234, 500),
-            TestNestedChildMessage("bar", bytearray(b"test2"), True, 3.1415, -302),
+            TestNestedChildMessage("bar", bytearray(b"test2"), True, 3.1415, 302),
             TestNestedChildMessage("baz", bytearray(b"test3"), False, 1e5, 300),
         )
         self.instance_filled_nested_bytes = bytes(self.instance_filled_nested)
         self.instance_filled_repeated = TestRepeatedMessage(
-            [
-                "test1",
-                "test2",
-                "test3",
-                "test4",
-                "test5",
-                "test6",
-                "test7",
-                "test8",
-                "test9",
-                "test10",
-            ],
-            [2, -100, 0, 500000, 600, -425678, 1000000000, -300, 1, -694214214466],
-            [True, False, False, False, True, True, False, True, False, False],
+            [f"test{i}" for i in range(1_000)],
+            [(i - 500) ** 3 for i in range(1_000)],
+            [i % 2 == 0 for i in range(1_000)],
         )
         self.instance_filled_repeated_bytes = bytes(self.instance_filled_repeated)
 
@@ -71,9 +60,9 @@ def time_overhead(self):
 
         @dataclass
         class Message(betterproto.Message):
-            foo: int = betterproto.uint32_field(0)
-            bar: str = betterproto.string_field(1)
-            baz: float = betterproto.float_field(2)
+            foo: int = betterproto.uint32_field(1)
+            bar: str = betterproto.string_field(2)
+            baz: float = betterproto.float_field(3)
 
     def time_instantiation(self):
         """Time instantiation"""

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -20,6 +20,7 @@ jinja2 = { version = ">=3.0.3", optional = true }
 python-dateutil = "^2.8"
 isort = {version = "^5.11.5", optional = true}
 typing-extensions = "^4.7.1"
+betterproto-rust-codec = { version = "0.1.0", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 asv = "^0.4.2"
@@ -48,6 +49,7 @@ protoc-gen-python_betterproto = "betterproto.plugin:main"
 
 [tool.poetry.extras]
 compiler = ["black", "isort", "jinja2"]
+rust-codec = ["betterproto-rust-codec"]
 
 
 # Dev workflow tasks

diff --git a/src/betterproto/__init__.py b/src/betterproto/__init__.py
@@ -1868,6 +1868,23 @@ def _validate_field_groups(cls, values):
 
 Message.__annotations__ = {}  # HACK to avoid typing.get_type_hints breaking :)
 
+# monkey patch (de-)serialization functions of class `Message`
+# with functions from `betterproto-rust-codec` if available
+try:
+    import betterproto_rust_codec
+
+    def __parse_patch(self: T, data: bytes) -> T:
+        betterproto_rust_codec.deserialize(self, data)
+        return self
+
+    def __bytes_patch(self) -> bytes:
+        return betterproto_rust_codec.serialize(self)
+
+    Message.parse = __parse_patch
+    Message.__bytes__ = __bytes_patch
+except ModuleNotFoundError:
+    pass
+
 
 def serialized_on_wire(message: Message) -> bool:
     """