Protobuf to dataclass converstion.

adaptix_protobuf.py

from dataclasses import dataclass
from pathlib import Path

import grpc
from google._upb._message import FieldDescriptor

from adaptix import TypeHint
from adaptix._internal.conversion.facade.func import get_converter
from adaptix._internal.model_tools.definitions import (
    InputShape, InputField, NoDefault, Param, ParamKind, OutputShape,
    OutputField, create_attr_accessor, FullShape, Shape, IntrospectionError
)
from adaptix._internal.model_tools.introspection.sqlalchemy import IdWrapper
from adaptix._internal.provider.shape_provider import ShapeProvider

# .proto:
# message SubRequest { int32 value = 1; }
# message MyRequest {
#   string name = 1;
#   SubRequest data = 2;
# }

code_dir = Path(__file__).parent / "my_grpc_service.proto"
myprotos, myservices = grpc.protos_and_services(
    str(code_dir.relative_to(Path.cwd())),
)

TYPES = {
    FieldDescriptor.TYPE_BOOL: bool,
    FieldDescriptor.TYPE_DOUBLE: float,
    FieldDescriptor.TYPE_FLOAT: float,
    FieldDescriptor.TYPE_INT64: int,
    FieldDescriptor.TYPE_UINT64: int,
    FieldDescriptor.TYPE_INT32: int,
    FieldDescriptor.TYPE_UINT32: int,
    FieldDescriptor.TYPE_SINT32: int,
    FieldDescriptor.TYPE_SINT64: int,
    FieldDescriptor.TYPE_FIXED32: int,
    FieldDescriptor.TYPE_FIXED64: int,
    FieldDescriptor.TYPE_SFIXED32: int,
    FieldDescriptor.TYPE_SFIXED64: int,
    FieldDescriptor.TYPE_STRING: str,
    FieldDescriptor.TYPE_BYTES: bytes,
}
MESSAGES = {
    t.DESCRIPTOR: t
    for t in vars(myprotos).values()
    if hasattr(t, "DESCRIPTOR")
}


def _get_field_type(desc: FieldDescriptor) -> TypeHint:
    if desc.type == FieldDescriptor.TYPE_MESSAGE:
        return MESSAGES[desc.message_type]
    return TYPES[desc.type]


def _get_default(desc: FieldDescriptor) -> TypeHint:
    return (
        desc.default_value
        if desc.has_default_value
        else NoDefault()
    )


def _get_input_shape(tp: TypeHint) -> InputShape:
    fields = []
    params = []
    for name, field_desc in tp.DESCRIPTOR.fields_by_name.items():
        fields.append(
            InputField(
                id=name,
                type=_get_field_type(field_desc),
                default=_get_default(field_desc),
                is_required=False,
                metadata={},
                original=IdWrapper(field_desc),
            ),
        )
        params.append(
            Param(
                field_id=name,
                name=name,
                kind=ParamKind.KW_ONLY,
            ),
        )
    return InputShape(
        constructor=tp,
        fields=tuple(fields),
        overriden_types=frozenset(),
        kwargs=None,
        params=tuple(params),
    )


def _get_output_shape(tp: TypeHint) -> OutputShape:
    output_fields = [
        OutputField(
            id=name,
            type=_get_field_type(field_desc),
            default=_get_default(field_desc),
            metadata={},
            original=IdWrapper(field_desc),
            accessor=create_attr_accessor(name, is_required=True),
        )
        for name, field_desc in tp.DESCRIPTOR.fields_by_name.items()
    ]
    return OutputShape(
        fields=tuple(output_fields),
        overriden_types=frozenset(),
    )


def get_protobuf_shape(tp) -> FullShape:
    if not hasattr(tp, "DESCRIPTOR"):
        raise IntrospectionError
    return Shape(
        input=_get_input_shape(tp=tp),
        output=_get_output_shape(tp=tp),
    )


shape_provider = ShapeProvider(get_protobuf_shape)


# usage:

@dataclass
class SubRequest:
    value: int


@dataclass
class MyRequest:
    name: str
    data: SubRequest


converter = get_converter(myprotos.MyRequest, MyRequest,
                          recipe=[shape_provider])

original = myprotos.MyRequest(
    name="hello",
    data=myprotos.SubRequest(value=42),
)
res = converter(original)
print(res)