Support binary serialization of ProtoBuf schemas

karapace/protobuf/enum_element.py

@@ -4,6 +4,8 @@
 """
 # Ported from square/wire:
 # wire-library/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/internal/parser/EnumElement.kt
+from __future__ import annotations
+
 from itertools import chain
 from karapace.protobuf.compare_result import CompareResult, Modification
 from karapace.protobuf.compare_type_storage import CompareTypes
@@ -12,7 +14,6 @@
 from karapace.protobuf.option_element import OptionElement
 from karapace.protobuf.type_element import TypeElement
 from karapace.protobuf.utils import append_documentation, append_indented
-from typing import List
 
 
 class EnumElement(TypeElement):
@@ -21,8 +22,8 @@ def __init__(
         location: Location,
         name: str,
         documentation: str = "",
-        options: List[OptionElement] = None,
-        constants: List[EnumConstantElement] = None,
+        options: list[OptionElement] | None = None,
+        constants: list[EnumConstantElement] | None = None,
     ) -> None:
         # Enums do not allow nested type declarations.
         super().__init__(location, name, documentation, options or [], [])
@@ -47,7 +48,7 @@ def to_schema(self) -> str:
         result.append("}\n")
         return "".join(result)
 
-    def compare(self, other: "EnumElement", result: CompareResult, types: CompareTypes) -> None:
+    def compare(self, other: EnumElement, result: CompareResult, types: CompareTypes) -> None:
         self_tags = {}
         other_tags = {}
         constant: EnumConstantElement

karapace/protobuf/location.py

@@ -48,3 +48,6 @@ def __str__(self) -> str:
                 result += str(self.column)
 
         return result
+
+
+DEFAULT_LOCATION = Location("", "")

karapace/protobuf/message_element.py

@@ -5,6 +5,8 @@
 # Ported from square/wire:
 # wire-library/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/internal/parser/MessageElement.kt
 # compatibility routine added
+from __future__ import annotations
+
 from itertools import chain
 from karapace.protobuf.compare_result import CompareResult, Modification
 from karapace.protobuf.compare_type_storage import CompareTypes
@@ -17,7 +19,6 @@
 from karapace.protobuf.reserved_element import ReservedElement
 from karapace.protobuf.type_element import TypeElement
 from karapace.protobuf.utils import append_documentation, append_indented
-from typing import List
 
 
 class MessageElement(TypeElement):
@@ -26,13 +27,13 @@ def __init__(
         location: Location,
         name: str,
         documentation: str = "",
-        nested_types: List[TypeElement] = None,
-        options: List[OptionElement] = None,
-        reserveds: List[ReservedElement] = None,
-        fields: List[FieldElement] = None,
-        one_ofs: List[OneOfElement] = None,
-        extensions: List[ExtensionsElement] = None,
-        groups: List[GroupElement] = None,
+        nested_types: list[TypeElement] | None = None,
+        options: list[OptionElement] | None = None,
+        reserveds: list[ReservedElement] | None = None,
+        fields: list[FieldElement] | None = None,
+        one_ofs: list[OneOfElement] | None = None,
+        extensions: list[ExtensionsElement] | None = None,
+        groups: list[GroupElement] | None = None,
     ) -> None:
         super().__init__(location, name, documentation, options or [], nested_types or [])
         self.reserveds = reserveds or []
@@ -42,7 +43,7 @@ def __init__(
         self.groups = groups or []
 
     def to_schema(self) -> str:
-        result = []
+        result: list[str] = []
         append_documentation(result, self.documentation)
         result.append(f"message {self.name} {{")
         if self.reserveds:
@@ -83,7 +84,7 @@ def to_schema(self) -> str:
         result.append("}\n")
         return "".join(result)
 
-    def compare(self, other: "MessageElement", result: CompareResult, types: CompareTypes) -> None:
+    def compare(self, other: MessageElement, result: CompareResult, types: CompareTypes) -> None:
         from karapace.protobuf.compare_type_lists import compare_type_lists
 
         if types.lock_message(self):

karapace/protobuf/schema.py

@@ -15,17 +15,19 @@
 from karapace.protobuf.enum_element import EnumElement
 from karapace.protobuf.exception import IllegalArgumentException
 from karapace.protobuf.known_dependency import DependenciesHardcoded, KnownDependency
-from karapace.protobuf.location import Location
+from karapace.protobuf.location import DEFAULT_LOCATION
 from karapace.protobuf.message_element import MessageElement
 from karapace.protobuf.one_of_element import OneOfElement
 from karapace.protobuf.option_element import OptionElement
 from karapace.protobuf.proto_file_element import ProtoFileElement
 from karapace.protobuf.proto_parser import ProtoParser
+from karapace.protobuf.serialization import deserialize, serialize
 from karapace.protobuf.type_element import TypeElement
 from karapace.protobuf.utils import append_documentation, append_indented
 from karapace.schema_references import Reference
 from typing import Iterable, Mapping, Sequence
 
+import binascii
 import itertools
 
 
@@ -247,19 +249,25 @@ def add_new_type(
 
 
 class ProtobufSchema:
-    DEFAULT_LOCATION = Location("", "")
-
     def __init__(
         self,
         schema: str,
         references: Sequence[Reference] | None = None,
         dependencies: Mapping[str, Dependency] | None = None,
+        proto_file_element: ProtoFileElement | None = None,
     ) -> None:
         if type(schema).__name__ != "str":
             raise IllegalArgumentException("Non str type of schema string")
-        self.dirty = schema
         self.cache_string = ""
-        self.proto_file_element = ProtoParser.parse(self.DEFAULT_LOCATION, schema)
+
+        if proto_file_element is not None:
+            self.proto_file_element = proto_file_element
+        else:
+            try:
+                self.proto_file_element = deserialize(schema)
+            except binascii.Error:  # If not base64 formatted
+                self.proto_file_element = ProtoParser.parse(DEFAULT_LOCATION, schema)
+
         self.references = references
         self.dependencies = dependencies
 
@@ -573,3 +581,6 @@ def compare(self, other: ProtobufSchema, result: CompareResult) -> CompareResult
             self_dependencies=self.dependencies,
             other_dependencies=other.dependencies,
         )
+
+    def serialize(self) -> str:
+        return serialize(self.proto_file_element)