from abc import ABC, abstractmethod
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, ClassVar
from ._core import ModelBase
@dataclass(frozen=True, order=True)
class SemVer:
"""
Conversion and sorting of Semantic Versions.
This is a very simple implementation to avoid having to rely on an external
package like pydantic-semver.
"""
major: int
minor: int
patch: int
@classmethod
def parse(cls, s: str):
return cls(*(int(x) for x in s.split(".")))
def __str__(self):
return f"{self.major}.{self.minor}.{self.patch}"
@dataclass(frozen=True)
class MigrationInfo:
"""Stores info about a migration between two versions."""
from_version: str
to_version: str
func: Callable
def apply(self, migration):
"""Apply the migration."""
self.func(migration)
[docs]
class Migration(ABC):
"""Abstract class for migrations."""
[docs]
@abstractmethod
def add(self, path: str):
"""Add a new field or set of fields to the model."""
pass
[docs]
@abstractmethod
def rename(self, old: str, new: str):
"""Rename a field from old to new."""
pass
[docs]
@abstractmethod
def drop(self, path: str):
"""Drop field."""
pass
[docs]
@abstractmethod
def set_default(self, path: str, value):
"""Change a field's default."""
pass
class DictMigration(Migration):
"""Implementation of Migration that modifies dicts."""
def __init__(self, data: dict):
self.data = data
def _parent(self, path) -> tuple[dict, str]:
*parents, key = path.split(".")
obj = self.data
for p in parents:
obj = obj[p]
return obj, key
def add(self, path: str):
pass # no op needed for migration
def rename(self, old, new):
try:
old_parent, old_key = self._parent(old)
new_parent, new_key = self._parent(new)
except KeyError:
return
if old_key in old_parent:
new_parent[new_key] = old_parent.pop(old_key)
def drop(self, path):
try:
parent, key = self._parent(path)
except KeyError:
return
if key in parent:
parent.pop(key)
def set_default(self, path, value):
try:
parent, key = self._parent(path)
except KeyError:
return
parent.setdefault(key, value)
def update_metadata(self, path: str, metadata_field: str, old: Any, new: Any):
pass # for now, the dict is assumed to be already translated from FITS, so no op is needed.
def apply_transform(self, path: str, func: Callable, description: str):
try:
parent, key = self._parent(path)
except KeyError:
return
func(obj=parent, key=key)
class HistoryMigration(Migration):
"""Implementation of Migration that records migration history."""
def __init__(self):
self.operations = []
def add(self, path: str):
self.operations.append(f"Added '{path}'")
def rename(self, old, new):
self.operations.append(f"Renamed '{old}' to '{new}'")
def drop(self, path):
self.operations.append(f"Dropped '{path}'")
def set_default(self, path, value):
self.operations.append(f"Changed default value of '{path}' = '{value!r}'")
def update_metadata(self, path: str, metadata_field: str, old: Any, new: Any):
self.operations.append(
f"Updated '{metadata_field}' metadata for '{path}' from '{old}' to '{new}'"
)
def apply_transform(self, path: str, func: Callable, description: str):
self.operations.append(
f"Transformed '{path}' using {func.__name__} ({description})"
)
class FITSHeaderMigration(Migration):
"""Migrations that are applied before FITS header conversion."""
def __init__(self, fits_header):
self.header = fits_header
def add(self, path: str):
pass
def rename(self, old, new):
pass
def drop(self, path):
pass
def set_default(self, path, value):
pass
def update_metadata(self, path: str, metadata_field: str, old: Any, new: Any):
if metadata_field == "fits_keyword":
if old in self.header:
self.header.rename_keyword(old, new)
def apply_transform(self, path: str, func: Callable, description: str):
pass
[docs]
def migration(from_version: str, to_version: str):
"""Migration decorator."""
def decorator(func):
func.__migration__ = MigrationInfo(
from_version=from_version,
to_version=to_version,
func=func,
)
return staticmethod(func)
return decorator
[docs]
class VersionedModel(ModelBase):
"""Base class for ModelBases that are versioned."""
_migrations: ClassVar[list[MigrationInfo]] #: internal storage of migration info
@classmethod
def __pydantic_init_subclass__(cls, **kwargs):
super().__pydantic_init_subclass__(**kwargs)
cls._migrations = [
member.__func__.__migration__
for _, member in cls.__dict__.items()
if isinstance(member, staticmethod)
and hasattr(member.__func__, "__migration__")
]
[docs]
@classmethod
def model_validate_versioned(cls, data: dict):
"""
Like model_validate, but applies version migrations.
Migrations are applied to the dict-like object to validate if needed,
before passing to pydantic for final validation.
"""
migration = DictMigration(dict(data))
for info in cls._migrations:
info.apply(migration)
return cls.model_validate(migration.data)
[docs]
@classmethod
def migration_history(cls):
"""Return the version migration history for this model."""
history = []
for info in cls._migrations:
recorder = HistoryMigration()
info.apply(recorder)
history.append(
{
"from": info.from_version,
"to": info.to_version,
"operations": recorder.operations,
}
)
return history
[docs]
@classmethod
def print_history(cls, printer=print):
"""Print a human-readable version of the history."""
printer(cls.__name__)
printer("=" * len(cls.__name__))
history = cls.migration_history()
for revision in history:
printer(f"* {revision['from']} -> {revision['to']}")
for op in revision["operations"]:
printer(f" - {op}")
