from functools import wraps

from marshmallow import Schema


def proto_schemas(input_schema=None, output_schema=None):
    """Decorator for serializing input arguments and deserializing returned value."""

    def decorator(func):
        @wraps(func)
        async def wrapper(data=None, **kwargs):
            if input_schema is not None:
                if data is None:
                    raise RuntimeError(
                        "This api_provider requires protobuf message input"
                    )
                dumped = input_schema().from_bytes(data)
                kwargs.update(dumped)

            got = await func(**kwargs)
            if output_schema is not None:
                return output_schema().to_bytes(got)

        return wrapper

    return decorator


class StrictSchema(Schema):
    def __init__(self, *args, **kwargs):
        kwargs["strict"] = True
        super().__init__(*args, **kwargs)


__all__ = ["proto_schemas", "StrictSchema"]
