betterproto2-compiler 0.2.0__py3-none-any.whl

betterproto2_compiler/plugin/module_validation.py

@@ -0,0 +1,156 @@
+import re
+from collections import defaultdict
+from collections.abc import Iterator
+from dataclasses import (
+    dataclass,
+    field,
+)
+
+
+@dataclass
+class ModuleValidator:
+    line_iterator: Iterator[str]
+    line_number: int = field(init=False, default=0)
+
+    collisions: dict[str, list[tuple[int, str]]] = field(init=False, default_factory=lambda: defaultdict(list))
+
+    def add_import(self, imp: str, number: int, full_line: str):
+        """
+        Adds an import to be tracked.
+        """
+        self.collisions[imp].append((number, full_line))
+
+    def process_import(self, imp: str):
+        """
+        Filters out the import to its actual value.
+        """
+        if " as " in imp:
+            imp = imp[imp.index(" as ") + 4 :]
+
+        imp = imp.strip()
+        assert " " not in imp, imp
+        return imp
+
+    def evaluate_multiline_import(self, line: str):
+        """
+        Evaluates a multiline import from a starting line
+        """
+        # Filter the first line and remove anything before the import statement.
+        full_line = line
+        line = line.split("import", 1)[1]
+        if "(" in line:
+            conditional = lambda line: ")" not in line
+        else:
+            conditional = lambda line: "\\" in line
+
+        # Remove open parenthesis if it exists.
+        if "(" in line:
+            line = line[line.index("(") + 1 :]
+
+        # Choose the conditional based on how multiline imports are formatted.
+        while conditional(line):
+            # Split the line by commas
+            imports = line.split(",")
+
+            for imp in imports:
+                # Add the import to the namespace
+                imp = self.process_import(imp)
+                if imp:
+                    self.add_import(imp, self.line_number, full_line)
+            # Get the next line
+            full_line = line = next(self.line_iterator)
+            # Increment the line number
+            self.line_number += 1
+
+        # validate the last line
+        if ")" in line:
+            line = line[: line.index(")")]
+        imports = line.split(",")
+        for imp in imports:
+            imp = self.process_import(imp)
+            if imp:
+                self.add_import(imp, self.line_number, full_line)
+
+    def evaluate_import(self, line: str):
+        """
+        Extracts an import from a line.
+        """
+        whole_line = line
+        line = line[line.index("import") + 6 :]
+        values = line.split(",")
+        for v in values:
+            self.add_import(self.process_import(v), self.line_number, whole_line)
+
+    def next(self):
+        """
+        Evaluate each line for names in the module.
+        """
+        line = next(self.line_iterator)
+
+        # Skip lines with indentation or comments
+        if (
+            # Skip indents and whitespace.
+            line.startswith(" ")
+            or line == "\n"
+            or line.startswith("\t")
+            or
+            # Skip comments
+            line.startswith("#")
+            or
+            # Skip  decorators
+            line.startswith("@")
+        ):
+            self.line_number += 1
+            return
+
+        # Skip docstrings.
+        if line.startswith('"""') or line.startswith("'''"):
+            quote = line[0] * 3
+            line = line[3:]
+            while quote not in line:
+                line = next(self.line_iterator)
+            self.line_number += 1
+            return
+
+        # Evaluate Imports.
+        if line.startswith("from ") or line.startswith("import "):
+            if "(" in line or "\\" in line:
+                self.evaluate_multiline_import(line)
+            else:
+                self.evaluate_import(line)
+
+        # Evaluate Classes.
+        elif match := re.search(r"^class (\w+)", line):
+            class_name = match.group(1)
+            if class_name:
+                self.add_import(class_name, self.line_number, line)
+
+        # Evaluate Functions.
+        elif match := re.search(r"^def (\w+)", line):
+            function_name = match.group(1)
+            if function_name:
+                self.add_import(function_name, self.line_number, line)
+
+        # Evaluate direct assignments.
+        elif match := re.search(r"^(\w+)\s*=", line):
+            assignment = match.group(1)
+            if assignment:
+                self.add_import(assignment, self.line_number, line)
+
+        self.line_number += 1
+
+    def validate(self) -> bool:
+        """
+        Run Validation.
+        """
+        try:
+            while True:
+                self.next()
+        except StopIteration:
+            pass
+
+        # Filter collisions for those with more than one value.
+        self.collisions = {k: v for k, v in self.collisions.items() if len(v) > 1}
+
+        # Return True if no collisions are found.
+        return not bool(self.collisions)

betterproto2_compiler/plugin/parser.py

@@ -0,0 +1,272 @@
+import pathlib
+import sys
+from collections.abc import Generator
+
+from betterproto2_compiler.lib.google.protobuf import (
+    DescriptorProto,
+    EnumDescriptorProto,
+    FileDescriptorProto,
+    ServiceDescriptorProto,
+)
+from betterproto2_compiler.lib.google.protobuf.compiler import (
+    CodeGeneratorRequest,
+    CodeGeneratorResponse,
+    CodeGeneratorResponseFeature,
+    CodeGeneratorResponseFile,
+)
+from betterproto2_compiler.settings import Settings
+
+from .compiler import outputfile_compiler
+from .models import (
+    EnumDefinitionCompiler,
+    FieldCompiler,
+    MapEntryCompiler,
+    MessageCompiler,
+    OneofCompiler,
+    OneOfFieldCompiler,
+    OutputTemplate,
+    PluginRequestCompiler,
+    ServiceCompiler,
+    ServiceMethodCompiler,
+    is_map,
+    is_oneof,
+)
+from .typing_compiler import (
+    DirectImportTypingCompiler,
+    NoTyping310TypingCompiler,
+    TypingImportTypingCompiler,
+)
+
+
+def traverse(
+    proto_file: FileDescriptorProto,
+) -> Generator[tuple[EnumDescriptorProto | DescriptorProto, list[int]], None, None]:
+    # Todo: Keep information about nested hierarchy
+    def _traverse(
+        path: list[int],
+        items: list[EnumDescriptorProto] | list[DescriptorProto],
+        prefix: str = "",
+    ) -> Generator[tuple[EnumDescriptorProto | DescriptorProto, list[int]], None, None]:
+        for i, item in enumerate(items):
+            # Adjust the name since we flatten the hierarchy.
+            # Todo: don't change the name, but include full name in returned tuple
+            should_rename = not isinstance(item, DescriptorProto) or not item.options or not item.options.map_entry
+
+            item.name = next_prefix = f"{prefix}.{item.name}" if prefix and should_rename else item.name
+            yield item, [*path, i]
+
+            if isinstance(item, DescriptorProto):
+                # Get nested types.
+                yield from _traverse([*path, i, 4], item.enum_type, next_prefix)
+                yield from _traverse([*path, i, 3], item.nested_type, next_prefix)
+
+    yield from _traverse([5], proto_file.enum_type)
+    yield from _traverse([4], proto_file.message_type)
+
+
+def get_settings(plugin_options: list[str]) -> Settings:
+    # Gather any typing generation options.
+    typing_opts = [opt[len("typing.") :] for opt in plugin_options if opt.startswith("typing.")]
+
+    if len(typing_opts) > 1:
+        raise ValueError("Multiple typing options provided")
+
+    # Set the compiler type.
+    typing_opt = typing_opts[0] if typing_opts else "direct"
+    if typing_opt == "direct":
+        typing_compiler = DirectImportTypingCompiler()
+    elif typing_opt == "root":
+        typing_compiler = TypingImportTypingCompiler()
+    elif typing_opt == "310":
+        typing_compiler = NoTyping310TypingCompiler()
+    else:
+        raise ValueError("Invalid typing option provided")
+
+    return Settings(
+        typing_compiler=typing_compiler,
+        pydantic_dataclasses="pydantic_dataclasses" in plugin_options,
+    )
+
+
+def generate_code(request: CodeGeneratorRequest) -> CodeGeneratorResponse:
+    response = CodeGeneratorResponse(supported_features=CodeGeneratorResponseFeature.FEATURE_PROTO3_OPTIONAL)
+
+    plugin_options = request.parameter.split(",") if request.parameter else []
+    settings = get_settings(plugin_options)
+
+    request_data = PluginRequestCompiler(plugin_request_obj=request)
+    # Gather output packages
+    for proto_file in request.proto_file:
+        output_package_name = proto_file.package
+        if output_package_name not in request_data.output_packages:
+            # Create a new output if there is no output for this package
+            request_data.output_packages[output_package_name] = OutputTemplate(
+                parent_request=request_data, package_proto_obj=proto_file, settings=settings
+            )
+        # Add this input file to the output corresponding to this package
+        request_data.output_packages[output_package_name].input_files.append(proto_file)
+
+    # Read Messages and Enums
+    # We need to read Messages before Services in so that we can
+    # get the references to input/output messages for each service
+    for output_package_name, output_package in request_data.output_packages.items():
+        for proto_input_file in output_package.input_files:
+            for item, path in traverse(proto_input_file):
+                read_protobuf_type(
+                    source_file=proto_input_file,
+                    item=item,
+                    path=path,
+                    output_package=output_package,
+                )
+
+    # Read Services
+    for output_package_name, output_package in request_data.output_packages.items():
+        for proto_input_file in output_package.input_files:
+            for index, service in enumerate(proto_input_file.service):
+                read_protobuf_service(proto_input_file, service, index, output_package)
+
+    # All the hierarchy is ready. We can perform pre-computations before generating the output files
+    for package in request_data.output_packages.values():
+        for message in package.messages.values():
+            for field in message.fields:
+                field.ready()
+            message.ready()
+        for enum in package.enums.values():
+            enum.ready()
+        for service in package.services.values():
+            for method in service.methods:
+                method.ready()
+            service.ready()
+
+    # Generate output files
+    output_paths: set[pathlib.Path] = set()
+    for output_package_name, output_package in request_data.output_packages.items():
+        # Add files to the response object
+        output_path = pathlib.Path(*output_package_name.split("."), "__init__.py")
+        output_paths.add(output_path)
+
+        response.file.append(
+            CodeGeneratorResponseFile(
+                name=str(output_path),
+                # Render and then format the output file
+                content=outputfile_compiler(output_file=output_package),
+            ),
+        )
+
+    # Make each output directory a package with __init__ file
+    init_files = {
+        directory.joinpath("__init__.py")
+        for path in output_paths
+        for directory in path.parents
+        if not directory.joinpath("__init__.py").exists()
+    } - output_paths
+
+    for init_file in init_files:
+        response.file.append(CodeGeneratorResponseFile(name=str(init_file)))
+
+    response.file.append(
+        CodeGeneratorResponseFile(
+            name="message_pool.py", content="import betterproto2\n\ndefault_message_pool = betterproto2.MessagePool()\n"
+        )
+    )
+
+    for output_package_name in sorted(output_paths.union(init_files)):
+        print(f"Writing {output_package_name}", file=sys.stderr)
+
+    return response
+
+
+def read_protobuf_type(
+    item: DescriptorProto | EnumDescriptorProto,
+    path: list[int],
+    source_file: "FileDescriptorProto",
+    output_package: OutputTemplate,
+) -> None:
+    if isinstance(item, DescriptorProto):
+        if item.options and item.options.map_entry:
+            # Skip generated map entry messages since we just use dicts
+            return
+        # Process Message
+        message_data = MessageCompiler(
+            source_file=source_file,
+            output_file=output_package,
+            proto_obj=item,
+            path=path,
+        )
+        output_package.messages[message_data.proto_name] = message_data
+
+        for index, field in enumerate(item.field):
+            if is_map(field, item):
+                message_data.fields.append(
+                    MapEntryCompiler(
+                        source_file=source_file,
+                        message=message_data,
+                        proto_obj=field,
+                        path=path + [2, index],
+                        typing_compiler=output_package.settings.typing_compiler,
+                    )
+                )
+            elif is_oneof(field):
+                message_data.fields.append(
+                    OneOfFieldCompiler(
+                        source_file=source_file,
+                        message=message_data,
+                        proto_obj=field,
+                        path=path + [2, index],
+                        typing_compiler=output_package.settings.typing_compiler,
+                    )
+                )
+            else:
+                message_data.fields.append(
+                    FieldCompiler(
+                        source_file=source_file,
+                        message=message_data,
+                        proto_obj=field,
+                        path=path + [2, index],
+                        typing_compiler=output_package.settings.typing_compiler,
+                    )
+                )
+
+        for index, oneof in enumerate(item.oneof_decl):
+            message_data.oneofs.append(
+                OneofCompiler(
+                    source_file=source_file,
+                    path=path + [8, index],
+                    proto_obj=oneof,
+                )
+            )
+
+    elif isinstance(item, EnumDescriptorProto):
+        # Enum
+        enum = EnumDefinitionCompiler(
+            source_file=source_file,
+            output_file=output_package,
+            proto_obj=item,
+            path=path,
+        )
+        output_package.enums[enum.proto_name] = enum
+
+
+def read_protobuf_service(
+    source_file: FileDescriptorProto,
+    service: ServiceDescriptorProto,
+    index: int,
+    output_package: OutputTemplate,
+) -> None:
+    service_data = ServiceCompiler(
+        source_file=source_file,
+        output_file=output_package,
+        proto_obj=service,
+        path=[6, index],
+    )
+    service_data.output_file.services[service_data.proto_name] = service_data
+
+    for j, method in enumerate(service.method):
+        service_data.methods.append(
+            ServiceMethodCompiler(
+                source_file=source_file,
+                parent=service_data,
+                proto_obj=method,
+                path=[6, index, 2, j],
+            )
+        )

betterproto2_compiler/plugin/plugin.bat

@@ -0,0 +1,2 @@
+@SET plugin_dir=%~dp0
+@python -m %plugin_dir% %*

betterproto2_compiler/plugin/typing_compiler.py

@@ -0,0 +1,163 @@
+import abc
+import builtins
+from collections import defaultdict
+from collections.abc import Iterator
+from dataclasses import (
+    dataclass,
+    field,
+)
+
+
+class TypingCompiler(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def optional(self, type_: str) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def list(self, type_: str) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def dict(self, key: str, value: str) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def union(self, *types: str) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def iterable(self, type_: str) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def async_iterable(self, type_: str) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def async_iterator(self, type_: str) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def imports(self) -> builtins.dict[str, set[str] | None]:
+        """
+        Returns either the direct import as a key with none as value, or a set of
+        values to import from the key.
+        """
+        raise NotImplementedError
+
+    def import_lines(self) -> Iterator:
+        imports = self.imports()
+        for key, value in imports.items():
+            if value is None:
+                yield f"import {key}"
+            else:
+                yield f"from {key} import ("
+                for v in sorted(value):
+                    yield f"    {v},"
+                yield ")"
+
+
+@dataclass
+class DirectImportTypingCompiler(TypingCompiler):
+    _imports: dict[str, set[str]] = field(default_factory=lambda: defaultdict(set))
+
+    def optional(self, type_: str) -> str:
+        self._imports["typing"].add("Optional")
+        return f"Optional[{type_}]"
+
+    def list(self, type_: str) -> str:
+        self._imports["typing"].add("List")
+        return f"List[{type_}]"
+
+    def dict(self, key: str, value: str) -> str:
+        self._imports["typing"].add("Dict")
+        return f"Dict[{key}, {value}]"
+
+    def union(self, *types: str) -> str:
+        self._imports["typing"].add("Union")
+        return f"Union[{', '.join(types)}]"
+
+    def iterable(self, type_: str) -> str:
+        self._imports["typing"].add("Iterable")
+        return f"Iterable[{type_}]"
+
+    def async_iterable(self, type_: str) -> str:
+        self._imports["typing"].add("AsyncIterable")
+        return f"AsyncIterable[{type_}]"
+
+    def async_iterator(self, type_: str) -> str:
+        self._imports["typing"].add("AsyncIterator")
+        return f"AsyncIterator[{type_}]"
+
+    def imports(self) -> builtins.dict[str, set[str] | None]:
+        return {k: v if v else None for k, v in self._imports.items()}
+
+
+@dataclass
+class TypingImportTypingCompiler(TypingCompiler):
+    _imported: bool = False
+
+    def optional(self, type_: str) -> str:
+        self._imported = True
+        return f"typing.Optional[{type_}]"
+
+    def list(self, type_: str) -> str:
+        self._imported = True
+        return f"typing.List[{type_}]"
+
+    def dict(self, key: str, value: str) -> str:
+        self._imported = True
+        return f"typing.Dict[{key}, {value}]"
+
+    def union(self, *types: str) -> str:
+        self._imported = True
+        return f"typing.Union[{', '.join(types)}]"
+
+    def iterable(self, type_: str) -> str:
+        self._imported = True
+        return f"typing.Iterable[{type_}]"
+
+    def async_iterable(self, type_: str) -> str:
+        self._imported = True
+        return f"typing.AsyncIterable[{type_}]"
+
+    def async_iterator(self, type_: str) -> str:
+        self._imported = True
+        return f"typing.AsyncIterator[{type_}]"
+
+    def imports(self) -> builtins.dict[str, set[str] | None]:
+        if self._imported:
+            return {"typing": None}
+        return {}
+
+
+@dataclass
+class NoTyping310TypingCompiler(TypingCompiler):
+    _imports: dict[str, set[str]] = field(default_factory=lambda: defaultdict(set))
+
+    def optional(self, type_: str) -> str:
+        return f"{type_} | None"
+
+    def list(self, type_: str) -> str:
+        return f"list[{type_}]"
+
+    def dict(self, key: str, value: str) -> str:
+        return f"dict[{key}, {value}]"
+
+    def union(self, *types: str) -> str:
+        return f"{' | '.join(types)}"
+
+    def iterable(self, type_: str) -> str:
+        self._imports["collections.abc"].add("Iterable")
+        return f"Iterable[{type_}]"
+
+    def async_iterable(self, type_: str) -> str:
+        self._imports["collections.abc"].add("AsyncIterable")
+        return f"AsyncIterable[{type_}]"
+
+    def async_iterator(self, type_: str) -> str:
+        self._imports["collections.abc"].add("AsyncIterator")
+        return f"AsyncIterator[{type_}]"
+
+    def imports(self) -> builtins.dict[str, set[str] | None]:
+        return {k: v if v else None for k, v in self._imports.items()}

betterproto2_compiler/settings.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+from .plugin.typing_compiler import TypingCompiler
+
+
+@dataclass
+class Settings:
+    pydantic_dataclasses: bool
+    typing_compiler: TypingCompiler

betterproto2_compiler/templates/header.py.j2

@@ -0,0 +1,59 @@
+{# All the imports needed for this file. The useless imports will be removed by Ruff. #}
+
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# sources: {{ ', '.join(output_file.input_filenames) }}
+# plugin: python-betterproto2
+# This file has been @generated
+
+__all__ = (
+    {% for _, enum in output_file.enums|dictsort(by="key") %}
+        "{{ enum.py_name }}",
+    {%- endfor -%}
+    {% for _, message in output_file.messages|dictsort(by="key") %}
+        "{{ message.py_name }}",
+    {%- endfor -%}
+    {% for _, service in output_file.services|dictsort(by="key") %}
+        "{{ service.py_name }}Stub",
+        "{{ service.py_name }}Base",
+    {%- endfor -%}
+)
+
+import builtins
+import datetime
+import warnings
+
+{% if output_file.settings.pydantic_dataclasses %}
+from pydantic.dataclasses import dataclass
+from pydantic import model_validator
+{%- else -%}
+from dataclasses import dataclass
+{% endif %}
+
+{% set typing_imports = output_file.settings.typing_compiler.imports() %}
+{% if typing_imports %}
+{% for line in output_file.settings.typing_compiler.import_lines() %}
+{{ line }}
+{% endfor %}
+{% endif %}
+
+import betterproto2
+{% if output_file.services %}
+from betterproto2.grpc.grpclib_server import ServiceBase
+import grpclib
+{% endif %}
+
+from typing import TYPE_CHECKING
+
+{# Import the message pool of the generated code. #}
+{% if output_file.package %}
+from {{ "." * output_file.package.count(".") }}..message_pool import default_message_pool
+{% else %}
+from .message_pool import default_message_pool
+{% endif %}
+
+if TYPE_CHECKING:
+    import grpclib.server
+    from betterproto2.grpc.grpclib_client import MetadataLike
+    from grpclib.metadata import Deadline
+
+betterproto2.check_compiler_version("{{ version }}")