Make ProtoFile a container type

src/proto_file.py

@@ -9,7 +9,7 @@
 from src.proto_extend import ProtoExtend
 from src.proto_import import ProtoImport
 from src.proto_message import ProtoMessage
-from src.proto_node import ParsedProtoNode, ProtoNode, ProtoNodeDiff
+from src.proto_node import ParsedProtoNode, ProtoContainerNode, ProtoNode, ProtoNodeDiff
 from src.proto_option import ProtoOption
 from src.proto_package import ProtoPackage
 from src.proto_service import ProtoService
@@ -21,13 +21,40 @@ class ParsedProtoFileNode(ParsedProtoNode):
     remaining_source: str
 
 
-class ProtoFile(ProtoNode):
-    def __init__(self, syntax: ProtoSyntax, nodes: list[ProtoNode], *args, **kwargs):
+class ProtoFileHeaderNode(ProtoNode):
+    def __init__(
+        self, header_nodes: list[ProtoNode], syntax: ProtoSyntax, *args, **kwargs
+    ):
         super().__init__(*args, **kwargs)
+        self.header_nodes = header_nodes
         self.syntax = syntax
-        self.nodes = nodes
 
-        if len([node for node in nodes if isinstance(node, ProtoPackage)]) > 1:
+    @classmethod
+    def match(
+        cls,
+        proto_source: str,
+        parent: Optional["ProtoNode"] = None,
+    ) -> Optional["ParsedProtoNode"]:
+        pass
+
+    def serialize(self) -> str:
+        raise NotImplementedError
+
+    def normalize(self) -> Optional["ProtoNode"]:
+        raise NotImplementedError
+
+
+class ParsedProtoFileHeaderNode(ParsedProtoNode):
+    node: "ProtoFileHeaderNode"
+    remaining_source: str
+
+
+class ProtoFile(ProtoContainerNode):
+    def __init__(self, syntax: ProtoSyntax, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.syntax = syntax
+
+        if len([node for node in self.nodes if isinstance(node, ProtoPackage)]) > 1:
             raise ValueError(f"Proto can't have more than one package statement")
 
     @property
@@ -53,9 +80,34 @@ def enums(self) -> list[ProtoEnum]:
     def messages(self) -> list[ProtoMessage]:
         return [node for node in self.nodes if isinstance(node, ProtoMessage)]
 
-    @staticmethod
-    def parse_partial_content(partial_proto_content: str) -> ParsedProtoNode:
-        node_types: list[type[ProtoNode]] = [
+    @classmethod
+    def match_header(
+        cls,
+        proto_source: str,
+        parent: Optional["ProtoNode"] = None,
+    ) -> Optional["ParsedProtoNode"]:
+        syntax, parsed_tree, remaining_source = cls.parse_syntax_and_preceding_comments(
+            proto_source.strip()
+        )
+        return ParsedProtoFileHeaderNode(
+            ProtoFileHeaderNode(list(parsed_tree), syntax, parent=parent),
+            remaining_source.strip(),
+        )
+
+    @classmethod
+    def match_footer(
+        cls,
+        proto_source: str,
+        parent: Optional[ProtoNode] = None,
+    ) -> Optional[str]:
+        trimmed_source = proto_source.strip()
+        if trimmed_source == "":
+            return ""
+        return None
+
+    @classmethod
+    def container_types(cls) -> list[type[ProtoNode]]:
+        return [
             ProtoImport,
             ProtoMessage,
             ProtoPackage,
@@ -66,16 +118,20 @@ def parse_partial_content(partial_proto_content: str) -> ParsedProtoNode:
             ProtoSingleLineComment,
             ProtoMultiLineComment,
         ]
-        for node_type in node_types:
-            try:
-                match_result = node_type.match(partial_proto_content)
-            except (ValueError, IndexError, TypeError):
-                raise ValueError(
-                    f"Could not parse proto content:\n{partial_proto_content}"
-                )
-            if match_result is not None:
-                return match_result
-        raise ValueError(f"Could not parse proto content:\n{partial_proto_content}")
+
+    @classmethod
+    def construct(
+        cls,
+        header_match: "ParsedProtoNode",
+        contained_nodes: list[ProtoNode],
+        footer_match: str,
+        parent: Optional[ProtoNode] = None,
+    ) -> ProtoNode:
+        assert isinstance(header_match, ParsedProtoFileHeaderNode)
+        return cls(
+            header_match.node.syntax,
+            header_match.node.header_nodes + contained_nodes,
+        )
 
     @staticmethod
     def parse_syntax_and_preceding_comments(
@@ -108,25 +164,6 @@ def parse_syntax_and_preceding_comments(
 
         return syntax, parsed_tree, proto_content
 
-    @classmethod
-    def match(
-        cls, proto_content: str, parent: Optional["ProtoNode"] = None
-    ) -> Optional[ParsedProtoFileNode]:
-        syntax, parsed_tree, proto_content = cls.parse_syntax_and_preceding_comments(
-            proto_content
-        )
-        new_tree: list[ProtoNode] = list(parsed_tree)
-        while proto_content:
-            # Remove empty statements.
-            if proto_content.startswith(";"):
-                proto_content = proto_content[1:].strip()
-                continue
-            match_result = cls.parse_partial_content(proto_content)
-            new_tree.append(match_result.node)
-            proto_content = match_result.remaining_source.strip()
-
-        return ParsedProtoFileNode(cls(syntax, new_tree), proto_content)
-
     def normalize(self) -> Optional["ProtoNode"]:
         normalized_nodes = [n.normalize() for n in self.nodes]
         return ProtoFile(

src/proto_node.py

@@ -98,6 +98,7 @@ def match(
 
         proto_source = header_match.remaining_source.strip()
         nodes = []
+        footer_match: Optional[str] = None
         while proto_source:
             # Remove empty statements.
             if proto_source.startswith(";"):
@@ -114,9 +115,11 @@ def match(
             proto_source = match_result.remaining_source.strip()
 
         if footer_match is None:
-            raise ValueError(
-                f"Footer was not found when matching container node {cls} for remaining proto source {proto_source}"
-            )
+            footer_match = cls.match_footer(proto_source, parent)
+            if footer_match is None:
+                raise ValueError(
+                    f"Footer was not found when matching container node {cls} for remaining proto source {proto_source}"
+                )
 
         return ParsedProtoNode(
             cls.construct(header_match, nodes, footer_match, parent=parent),

src/util/parser.py

@@ -13,12 +13,11 @@ def loads(proto_content: str) -> ProtoFile:
         try:
             parsed_file = ProtoFile.match(proto_content, None)
         except ValueError as e:
-            raise ParseError(
-                f"Proto doesn't have parseable syntax:\n{proto_content}\n{e}"
-            )
+            raise ParseError(f"Proto doesn't have parseable syntax:\n{e}")
         if parsed_file is None:
             raise ParseError(f"Proto doesn't have parseable syntax:\n{proto_content}")
 
+        assert isinstance(parsed_file.node, ProtoFile)
         return parsed_file.node