"""Schema for common dataloader schemas."""

from abacus_common_logic.marshalling.custom_fields import ma
from marshmallow import validates_schema


class DataloaderOutputSchema(ma.Schema):
    @validates_schema
    def validate_data_or_error(self, data, **kwargs):
        if 'data' not in data and 'error' not in data:
            raise ma.exceptions.ValidationError(
                "Each entry must have either 'data' or 'error'."
            )
        if 'data' in data and 'error' in data:
            raise ma.exceptions.ValidationError(
                "Each entry must have only one of 'data' or 'error'."
            )