Skip to content

Commit 4f18ed1

Browse files
Add support for streaming delimited messages (#529)
* Add support for streaming delimited messages This allows developers to easily dump and load multiple messages from a stream in a way that is compatible with official protobuf implementations (such as Java's `MessageLite#writeDelimitedTo(...)`). * Add Java compatibility tests for streaming These tests stream data such as messages to output files, have a Java binary read them and then write them back using the `protobuf-java` functions, and then read them back in on the Python side to check that the returned data is as expected. This checks that the official Java implementation (and so any other matching implementations) can properly parse outputs from Betterproto, and vice-versa, ensuring compatibility in these functions between the two. * Replace `xxxxableBuffer` with `SupportsXxxx`
1 parent 6b36b9b commit 4f18ed1

File tree

10 files changed

+537
-9
lines changed

10 files changed

+537
-9
lines changed

.pre-commit-config.yaml

+8-1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,13 @@ repos:
1616
- repo: https://github.com/PyCQA/doc8
1717
rev: 0.10.1
1818
hooks:
19-
- id: doc8
19+
- id: doc8
2020
additional_dependencies:
2121
- toml
22+
23+
- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
24+
rev: v2.10.0
25+
hooks:
26+
- id: pretty-format-java
27+
args: [--autofix, --aosp]
28+
files: ^.*\.java$

src/betterproto/__init__.py

+26-7
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,10 @@
5050

5151

5252
if TYPE_CHECKING:
53-
from _typeshed import ReadableBuffer
53+
from _typeshed import (
54+
SupportsRead,
55+
SupportsWrite,
56+
)
5457

5558

5659
# Proto 3 data types
@@ -127,6 +130,9 @@
127130
WIRE_FIXED_64_TYPES = [TYPE_DOUBLE, TYPE_FIXED64, TYPE_SFIXED64]
128131
WIRE_LEN_DELIM_TYPES = [TYPE_STRING, TYPE_BYTES, TYPE_MESSAGE, TYPE_MAP]
129132

133+
# Indicator of message delimitation in streams
134+
SIZE_DELIMITED = -1
135+
130136

131137
# Protobuf datetimes start at the Unix Epoch in 1970 in UTC.
132138
def datetime_default_gen() -> datetime:
@@ -322,7 +328,7 @@ def _pack_fmt(proto_type: str) -> str:
322328
}[proto_type]
323329

324330

325-
def dump_varint(value: int, stream: BinaryIO) -> None:
331+
def dump_varint(value: int, stream: "SupportsWrite[bytes]") -> None:
326332
"""Encodes a single varint and dumps it into the provided stream."""
327333
if value < -(1 << 63):
328334
raise ValueError(
@@ -531,7 +537,7 @@ def _dump_float(value: float) -> Union[float, str]:
531537
return value
532538

533539

534-
def load_varint(stream: BinaryIO) -> Tuple[int, bytes]:
540+
def load_varint(stream: "SupportsRead[bytes]") -> Tuple[int, bytes]:
535541
"""
536542
Load a single varint value from a stream. Returns the value and the raw bytes read.
537543
"""
@@ -569,7 +575,7 @@ class ParsedField:
569575
raw: bytes
570576

571577

572-
def load_fields(stream: BinaryIO) -> Generator[ParsedField, None, None]:
578+
def load_fields(stream: "SupportsRead[bytes]") -> Generator[ParsedField, None, None]:
573579
while True:
574580
try:
575581
num_wire, raw = load_varint(stream)
@@ -881,15 +887,19 @@ def _betterproto(self) -> ProtoClassMetadata:
881887
self.__class__._betterproto_meta = meta # type: ignore
882888
return meta
883889

884-
def dump(self, stream: BinaryIO) -> None:
890+
def dump(self, stream: "SupportsWrite[bytes]", delimit: bool = False) -> None:
885891
"""
886892
Dumps the binary encoded Protobuf message to the stream.
887893
888894
Parameters
889895
-----------
890896
stream: :class:`BinaryIO`
891897
The stream to dump the message to.
898+
delimit:
899+
Whether to prefix the message with a varint declaring its size.
892900
"""
901+
if delimit == SIZE_DELIMITED:
902+
dump_varint(len(self), stream)
893903

894904
for field_name, meta in self._betterproto.meta_by_field_name.items():
895905
try:
@@ -1207,7 +1217,11 @@ def _include_default_value_for_oneof(
12071217
meta.group is not None and self._group_current.get(meta.group) == field_name
12081218
)
12091219

1210-
def load(self: T, stream: BinaryIO, size: Optional[int] = None) -> T:
1220+
def load(
1221+
self: T,
1222+
stream: "SupportsRead[bytes]",
1223+
size: Optional[int] = None,
1224+
) -> T:
12111225
"""
12121226
Load the binary encoded Protobuf from a stream into this message instance. This
12131227
returns the instance itself and is therefore assignable and chainable.
@@ -1219,12 +1233,17 @@ def load(self: T, stream: BinaryIO, size: Optional[int] = None) -> T:
12191233
size: :class:`Optional[int]`
12201234
The size of the message in the stream.
12211235
Reads stream until EOF if ``None`` is given.
1236+
Reads based on a size delimiter prefix varint if SIZE_DELIMITED is given.
12221237
12231238
Returns
12241239
--------
12251240
:class:`Message`
12261241
The initialized message.
12271242
"""
1243+
# If the message is delimited, parse the message delimiter
1244+
if size == SIZE_DELIMITED:
1245+
size, _ = load_varint(stream)
1246+
12281247
# Got some data over the wire
12291248
self._serialized_on_wire = True
12301249
proto_meta = self._betterproto
@@ -1297,7 +1316,7 @@ def load(self: T, stream: BinaryIO, size: Optional[int] = None) -> T:
12971316

12981317
return self
12991318

1300-
def parse(self: T, data: "ReadableBuffer") -> T:
1319+
def parse(self: T, data: bytes) -> T:
13011320
"""
13021321
Parse the binary encoded Protobuf into this message instance. This
13031322
returns the instance itself and is therefore assignable and chainable.

tests/streams/delimited_messages.in

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
���:bTesting���:bTesting
2+
 

tests/streams/java/.gitignore

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
### Output ###
2+
target/
3+
!.mvn/wrapper/maven-wrapper.jar
4+
!**/src/main/**/target/
5+
!**/src/test/**/target/
6+
dependency-reduced-pom.xml
7+
MANIFEST.MF
8+
9+
### IntelliJ IDEA ###
10+
.idea/
11+
*.iws
12+
*.iml
13+
*.ipr
14+
15+
### Eclipse ###
16+
.apt_generated
17+
.classpath
18+
.factorypath
19+
.project
20+
.settings
21+
.springBeans
22+
.sts4-cache
23+
24+
### NetBeans ###
25+
/nbproject/private/
26+
/nbbuild/
27+
/dist/
28+
/nbdist/
29+
/.nb-gradle/
30+
build/
31+
!**/src/main/**/build/
32+
!**/src/test/**/build/
33+
34+
### VS Code ###
35+
.vscode/
36+
37+
### Mac OS ###
38+
.DS_Store

tests/streams/java/pom.xml

+94
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<project xmlns="http://maven.apache.org/POM/4.0.0"
3+
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
4+
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
5+
<modelVersion>4.0.0</modelVersion>
6+
7+
<groupId>betterproto</groupId>
8+
<artifactId>compatibility-test</artifactId>
9+
<version>1.0-SNAPSHOT</version>
10+
<packaging>jar</packaging>
11+
12+
<properties>
13+
<maven.compiler.source>11</maven.compiler.source>
14+
<maven.compiler.target>11</maven.compiler.target>
15+
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
16+
<protobuf.version>3.23.4</protobuf.version>
17+
</properties>
18+
19+
<dependencies>
20+
<dependency>
21+
<groupId>com.google.protobuf</groupId>
22+
<artifactId>protobuf-java</artifactId>
23+
<version>${protobuf.version}</version>
24+
</dependency>
25+
</dependencies>
26+
27+
<build>
28+
<extensions>
29+
<extension>
30+
<groupId>kr.motd.maven</groupId>
31+
<artifactId>os-maven-plugin</artifactId>
32+
<version>1.7.1</version>
33+
</extension>
34+
</extensions>
35+
36+
<plugins>
37+
<plugin>
38+
<groupId>org.apache.maven.plugins</groupId>
39+
<artifactId>maven-shade-plugin</artifactId>
40+
<version>3.5.0</version>
41+
<executions>
42+
<execution>
43+
<phase>package</phase>
44+
<goals>
45+
<goal>shade</goal>
46+
</goals>
47+
<configuration>
48+
<transformers>
49+
<transformer implementation="org.apache.maven.plugins.shade.resource.ManifestResourceTransformer">
50+
<mainClass>betterproto.CompatibilityTest</mainClass>
51+
</transformer>
52+
</transformers>
53+
</configuration>
54+
</execution>
55+
</executions>
56+
</plugin>
57+
58+
<plugin>
59+
<groupId>org.apache.maven.plugins</groupId>
60+
<artifactId>maven-jar-plugin</artifactId>
61+
<version>3.3.0</version>
62+
<configuration>
63+
<archive>
64+
<manifest>
65+
<addClasspath>true</addClasspath>
66+
<mainClass>betterproto.CompatibilityTest</mainClass>
67+
</manifest>
68+
</archive>
69+
</configuration>
70+
</plugin>
71+
72+
<plugin>
73+
<groupId>org.xolstice.maven.plugins</groupId>
74+
<artifactId>protobuf-maven-plugin</artifactId>
75+
<version>0.6.1</version>
76+
<executions>
77+
<execution>
78+
<goals>
79+
<goal>compile</goal>
80+
</goals>
81+
</execution>
82+
</executions>
83+
<configuration>
84+
<protocArtifact>
85+
com.google.protobuf:protoc:${protobuf.version}:exe:${os.detected.classifier}
86+
</protocArtifact>
87+
</configuration>
88+
</plugin>
89+
</plugins>
90+
91+
<finalName>${project.artifactId}</finalName>
92+
</build>
93+
94+
</project>
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package betterproto;
2+
3+
import java.io.IOException;
4+
5+
public class CompatibilityTest {
6+
public static void main(String[] args) throws IOException {
7+
if (args.length < 2)
8+
throw new RuntimeException("Attempted to run without the required arguments.");
9+
else if (args.length > 2)
10+
throw new RuntimeException(
11+
"Attempted to run with more than the expected number of arguments (>1).");
12+
13+
Tests tests = new Tests(args[1]);
14+
15+
switch (args[0]) {
16+
case "single_varint":
17+
tests.testSingleVarint();
18+
break;
19+
20+
case "multiple_varints":
21+
tests.testMultipleVarints();
22+
break;
23+
24+
case "single_message":
25+
tests.testSingleMessage();
26+
break;
27+
28+
case "multiple_messages":
29+
tests.testMultipleMessages();
30+
break;
31+
32+
case "infinite_messages":
33+
tests.testInfiniteMessages();
34+
break;
35+
36+
default:
37+
throw new RuntimeException(
38+
"Attempted to run with unknown argument '" + args[0] + "'.");
39+
}
40+
}
41+
}

0 commit comments

Comments
 (0)