import typing

from marshmallow import fields


class SaveToContextMixin:
    def __init__(self, *args, ctx_key=None, raw=False, **kwargs):
        self._ctx_key = ctx_key
        self._raw = raw
        super().__init__(*args, **kwargs)

    def _serialize(self, value: typing.Any, attr: str, obj: typing.Any, **kwargs):
        result = super()._serialize(value, attr, obj, **kwargs)
        self.parent.context[self._ctx_key or self.name] = value if self._raw else result
        return result


def save_ctx(field_cls: typing.Type[fields.Field], *args, **kwargs) -> fields.Field:
    class F(SaveToContextMixin, field_cls):
        pass

    return F(*args, **kwargs)