Sketch: Automatically migrate old data with alternatives from pydantic's BaseModels

auto_migrate.py

from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Self, override
from collections.abc import Callable
import pydantic


class FooV1(pydantic.BaseModel):
    a: int
    b: str
    c: float


class FooV2(pydantic.BaseModel):
    a: int
    b: str
    c: str
    d: float


class FooV3(pydantic.BaseModel):
    c: str


class MigrationError(BaseException):
    def __init__(self: Self, msg: str) -> None:
        super().__init__()
        self._msg: str = msg

    @override
    def __str__(self: Self) -> str:
        return self._msg


class Migration[In, Out](ABC):
    @abstractmethod
    def parse(self: Self, input: In) -> Out | MigrationError:
        raise NotImplementedError()

    def __call__(self: Self, input: In) -> Out | MigrationError:
        return self.parse(input)

    def parse_unsafe(self: Self, input: In) -> Out:
        result: Out | MigrationError = self.parse(input=input)
        if isinstance(result, MigrationError):
            raise result
        else:
            return result

    def with_alternative[T](
        self: Self, alternative: Migration[In, T]
    ) -> Migration[In, Out | T]:
        def _parse_fun(input: In) -> MigrationError | Out | T:
            result_a: Out | MigrationError = self.parse(input)
            if isinstance(result_a, MigrationError):
                result_b: T | MigrationError = alternative.parse(input)
                if isinstance(result_b, MigrationError):
                    return MigrationError(msg=f"{str(result_a)}\n{str(result_b)}")
                else:
                    return result_b
            else:
                return result_a

        return MigrationFunction(parse_fun=_parse_fun)

    def with_migration_alternative[T](
        self: Self,
        alternative: Migration[In, T],
        migration: Migration[Out, T],
    ) -> Migration[In, T]:
        return alternative.with_alternative(alternative=self.into(migration))

    def into[T](self: Self, migration: Migration[Out, T]) -> Migration[In, T]:
        def _parse_fun(input: In) -> T | MigrationError:
            parse_result: Out | MigrationError = self.parse(input)
            if isinstance(parse_result, MigrationError):
                return parse_result
            else:
                return migration(input=parse_result)

        return MigrationFunction(parse_fun=_parse_fun)

    def after[T](self: Self, migration: Migration[T, In]) -> Migration[T, Out]:
        def _parse_fun(input: T) -> Out | MigrationError:
            parse_result: In | MigrationError = migration(input)
            if isinstance(parse_result, MigrationError):
                return parse_result
            else:
                return self(input=parse_result)

        return MigrationFunction(parse_fun=_parse_fun)


class MigrationFunction[In, T](Migration[In, T]):
    def __init__(self: Self, parse_fun: Callable[[In], T | MigrationError]) -> None:
        self._parse_fun: Callable[[In], T | MigrationError] = parse_fun

    @override
    def parse(self: Self, input: In) -> T | MigrationError:
        try:
            return self._parse_fun(input)
        except Exception as exc:
            # TODO: Parse this into something better formatted
            return MigrationError(msg=str(exc))


class PydanticMigration[T: pydantic.BaseModel](Migration[str, T]):
    def __init__(self: Self, base_class: type[T]) -> None:
        self._base_class: type[T] = base_class

    @override
    def parse(self: Self, input: str) -> T | MigrationError:
        try:
            return self._base_class.model_validate_json(json_data=input)
        except pydantic.ValidationError as err:
            # TODO: Parse this into something better formatted
            msg = "\n".join(str(err) for err in err.errors(include_url=False))
            return MigrationError(msg=msg)


fooV1: Migration[str, FooV1] = PydanticMigration(base_class=FooV1)
fooV2: Migration[str, FooV2] = PydanticMigration(base_class=FooV2)
fooV3: Migration[str, FooV3] = PydanticMigration(base_class=FooV3)


def v1_to_v2(v: FooV1) -> FooV2:
    return FooV2(a=v.a, b=v.b, c="default", d=v.c)


def v2_to_v3(v: FooV2) -> FooV3:
    return FooV3(c=v.c)


migrate_to_fooV3: Migration[str, FooV3] = fooV1.with_migration_alternative(
    alternative=fooV2, migration=MigrationFunction(parse_fun=v1_to_v2)
).with_migration_alternative(
    alternative=fooV3, migration=MigrationFunction(parse_fun=v2_to_v3)
)