wip

karapace/protobuf/schema.py

@@ -2,6 +2,9 @@
 Copyright (c) 2023 Aiven Ltd
 See LICENSE for details
 """
+import itertools
+
+from karapace.dataclasses import default_dataclass
 # Ported from square/wire:
 # wire-library/wire-schema/src/commonMain/kotlin/com/squareup/wire/schema/Schema.kt
 # Ported partially for required functionality.
@@ -109,47 +112,106 @@ def option_element_string(option: OptionElement) -> str:
     return f"option {result};\n"
 
 
-@dataclasses.dataclass(frozen=True)
+@default_dataclass
 class UsedType:
     parent_object_qualified_name: str
     used_attribute_type: str
 
 
-@dataclasses.dataclass(frozen=True)
+@default_dataclass
+class SourceFileReference:
+    reference: str
+    import_order: int
+
+
+@default_dataclass
 class TypeTree:
     token: str
     children: List["TypeTree"]
+    original_file: Optional[SourceFileReference]
+
+    def expand(self) -> Optional[SourceFileReference]:
+        expanded_files = itertools.chain(
+            [] if self.original_file is None else [self.original_file],
+            (child.expand() for child in self.children)
+        )
+
+        # this is not empty by definition, you must have an import for each leaf at least
+        only_files = filter(
+            lambda x: x is None,
+            expanded_files
+        ),
+
+        return min(
+            only_files,
+            key=lambda x: x.imported_order
+        )
 
     @property
     def is_fully_qualified_type(self) -> bool:
+        return self.original_file is not None
+
+    @property
+    def is_leaf(self) -> bool:
         return len(self.children) == 0
 
     def represent(self, level=0) -> str:
         spacing = " " * 3 * level
-        if self.is_full_path_type:
-            return f"{spacing}{self.token}"
+        token = self.token + ">" if not self.is_fully_qualified_type else self.token
+        if self.is_leaf:
+            return f"{spacing}{token}"
         child_repr = "\n".join(child.represent(level=level + 1) for child in self.children)
-        return f"{spacing}{self.token} -> \n{child_repr}"
+        return f"{spacing}{token} -> \n{child_repr}"
 
     def __repr__(self) -> str:
         return self.represent()
 
 
-def add_new_type_r(parent_tree: TypeTree, remaining_tokens: List[str]) -> None:
+def _add_new_type(
+    parent_tree: TypeTree,
+    remaining_tokens: List[str],
+    file: str,
+    inserted_elements: int,
+) -> None:
     if remaining_tokens:
         token = remaining_tokens.pop()
         for child in parent_tree.children:
             if child.token == token:
-                return add_new_type_r(child, remaining_tokens)  # add a reference from which object/file was coming from
-        new_leaf = TypeTree(token, [])
+                return _add_new_type(
+                    child,
+                    remaining_tokens,
+                    file,
+                    inserted_elements
+                )  # add a reference from which object/file was coming from
+
+        new_leaf = TypeTree(
+            token=token,
+            children=[],
+            original_file=(
+                SourceFileReference(
+                    reference=file,
+                    import_order=inserted_elements
+                ) if remaining_tokens
+                else None
+            ),
+        )
         parent_tree.children.append(new_leaf)
-        return add_new_type_r(new_leaf, remaining_tokens)
-    # required by pylint
+        return _add_new_type(new_leaf, remaining_tokens, file, inserted_elements)
     return None
 
 
-def add_new_type(root_tree: TypeTree, full_path_type: str) -> None:
-    add_new_type_r(root_tree, full_path_type.split("."))
+def add_new_type(
+    root_tree: TypeTree,
+    full_path_type: str,
+    file: str,
+    inserted_elements: int,
+) -> None:
+    _add_new_type(
+        root_tree,
+        full_path_type.split("."),
+        file,
+        inserted_elements,
+    )  # one more it's added after that instruction
 
 
 class ProtobufSchema:
@@ -253,17 +315,29 @@ def verify_schema_dependencies(self) -> DependencyVerifierResult:
             return DependencyVerifierResult(True)
         return DependencyVerifierResult(False, f'type "{maybe_wrong_declaration}" is not defined')
 
-    def nested_type_tree(self, root_tree: TypeTree, parent_name: str, nested_type: TypeElement):
+    def nested_type_tree(
+        self,
+        root_tree: TypeTree,
+        parent_name: str,
+        nested_type: TypeElement,
+        filename: str,
+        inserted_types: int,
+    ) -> int:
         nested_component_full_path_name = parent_name + "." + nested_type.name
-        add_new_type(root_tree, nested_component_full_path_name)
+        add_new_type(root_tree, nested_component_full_path_name, filename, inserted_types)
+        inserted_types += 1
         for child in nested_type.nested_types:
-            self.nested_type_tree(root_tree, nested_component_full_path_name, child)
+            self.nested_type_tree(root_tree, nested_component_full_path_name, child, filename, inserted_types)
+            inserted_types += 1
+
+        return inserted_types
 
-    def types_tree_r(self, root_tree: TypeTree):
+    def _types_tree_recursive(self, root_tree: TypeTree, inserted_types: int, filename: str) -> tuple[TypeTree, int]:
         # verify that the import it's the same as the order of importing
         if self.dependencies:
-            for key in self.dependencies:
-                self.dependencies[key].get_schema().schema.types_tree_r(root_tree)
+            for dependency in self.dependencies:
+                dependency_tree, inserted_types = self.dependencies[dependency].get_schema().schema \
+                    ._types_tree_recursive(root_tree, inserted_types, dependency)
 
         # we can add an incremental number and a reference to the file
         # to get back which is the file who a certain declaration it's referring to
@@ -276,16 +350,21 @@ def types_tree_r(self, root_tree: TypeTree):
         for element_type in self.proto_file_element.types:
             type_name = element_type.name
             full_name = package_name + "." + type_name
-            add_new_type(root_tree, full_name)
+            add_new_type(root_tree, full_name, filename, inserted_types)
+            inserted_types += 1
 
             for nested_type in element_type.nested_types:
-                self.nested_type_tree(root_tree, full_name, nested_type)
+                inserted_types = self.nested_type_tree(root_tree, full_name, nested_type, filename, inserted_types)
 
-        return root_tree
+        return root_tree, inserted_types
 
     def types_tree(self) -> TypeTree:
-        root_tree = TypeTree(".", [])
-        self.types_tree_r(root_tree)
+        root_tree = TypeTree(
+            token=".",
+            children=[],
+            original_file=None,
+        )
+        self._types_tree_recursive(root_tree, 0, str(self.proto_file_element.location))
         return root_tree
 
     @staticmethod
@@ -294,17 +373,35 @@ def used_type(parent: str, element_type: str) -> List[UsedType]:
             end = element_type.find(">")
             virgule = element_type.find(",")
             key_element_type = element_type[4:virgule]
-            value_element_type = element_type[virgule + 1 : end]
+            value_element_type = element_type[virgule + 1: end]
             value_element_type = value_element_type.strip()
-            return [UsedType(parent, key_element_type), UsedType(parent, value_element_type)]
-        return [UsedType(parent, element_type)]
+            return [
+                UsedType(
+                    parent_object_qualified_name=parent,
+                    used_attribute_type=key_element_type
+                ), UsedType(
+                    parent_object_qualified_name=parent,
+                    used_attribute_type=value_element_type
+                )
+            ]
+        return [UsedType(
+            parent_object_qualified_name=parent,
+            used_attribute_type=element_type
+        )]
 
     @staticmethod
-    def dependencies_one_of(package_name: str, parent_name: str, one_of: OneOfElement) -> List[UsedType]:
+    def dependencies_one_of(
+        package_name: str,
+        parent_name: str,
+        one_of: OneOfElement,
+    ) -> List[UsedType]:
         parent = package_name + "." + parent_name
         dependencies = []
         for field in one_of.fields:
-            dependencies.append(UsedType(parent, field.element_type))
+            dependencies.append(UsedType(
+                parent_object_qualified_name=parent,
+                used_attribute_type=field.element_type,
+            ))
         return dependencies
 
     def used_types(self) -> List[UsedType]:
@@ -333,7 +430,12 @@ def used_types(self) -> List[UsedType]:
 
         return used_types
 
-    def nested_used_type(self, package_name: str, parent_name, element_type: TypeElement) -> List[str]:
+    def nested_used_type(
+        self,
+        package_name: str,
+        parent_name: str,
+        element_type: TypeElement,
+    ) -> List[str]:
         used_types = []
 
         if isinstance(element_type, MessageElement):

tests/unit/protobuf/test_protobuf_schema.py

@@ -406,6 +406,6 @@ def test_partial_path_in_protobuf() -> None:
     """
 
     no_ref_schema = ValidatedTypedSchema.parse(SchemaType.PROTOBUF, plan)
-    dep = Dependency("CustomerPlan.proto", Subject("foo"), 1, no_ref_schema)
-    ref_schema = ValidatedTypedSchema.parse(SchemaType.PROTOBUF, customer_plan_event, None, {"CustomerPlan.proto": dep})
+    dep = Dependency("ggwp", Subject("bar"), 1, no_ref_schema)
+    ref_schema = ValidatedTypedSchema.parse(SchemaType.PROTOBUF, customer_plan_event, None, {"foobar": dep})
     print(ref_schema)