optionally use betterproto-rust-codec

danielgtaylor · Dec 4, 2023 · aa6b9c5 · aa6b9c5
1 parent bd7de20
commit aa6b9c5
Show file tree

Hide file tree

Showing 4 changed files with 60 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
@@ -1014,6 +1014,13 @@ def __bytes__(self) -> bytes:
         """
         Get the binary encoded Protobuf representation of this message instance.
         """
+
+        try:
+            import betterproto_rust_codec
+            return betterproto_rust_codec.serialize(self)
+        except ModuleNotFoundError:
+            pass
+
         with BytesIO() as stream:
             self.dump(stream)
             return stream.getvalue()
@@ -1356,6 +1363,14 @@ def parse(self: T, data: bytes) -> T:
         :class:`Message`
             The initialized message.
         """
+
+        try:
+            import betterproto_rust_codec
+            betterproto_rust_codec.deserialize(self, data)
+            return self
+        except ModuleNotFoundError:
+            pass
+
         with BytesIO(data) as stream:
             return self.load(stream)