mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
tidy(config): use dataclass for MigrationEntry
The only pydantic usage was to convert strings to `Version` objects. The reason to do this conversion was to allow the register decorator to accept strings. MigrationEntry is only created inside this class, so we can just create versions from each migration when instantiating MigrationEntry instead. Also, pydantic doesn't provide runtime time checking for arbitrary classes like Version, so we don't get any real benefit.
This commit is contained in:
parent
aca9e44a3a
commit
6f128c86b4
@ -4,31 +4,22 @@
|
||||
Utility class for migrating among versions of the InvokeAI app config schema.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, List, TypeAlias
|
||||
|
||||
from packaging.version import Version
|
||||
from pydantic import BaseModel, ConfigDict, field_validator
|
||||
|
||||
AppConfigDict: TypeAlias = dict[str, Any]
|
||||
|
||||
|
||||
class MigrationEntry(BaseModel):
|
||||
@dataclass
|
||||
class MigrationEntry:
|
||||
"""Defines an individual migration."""
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
from_version: Version
|
||||
to_version: Version
|
||||
function: Callable[[AppConfigDict], AppConfigDict]
|
||||
|
||||
@field_validator("from_version", "to_version", mode="before")
|
||||
@classmethod
|
||||
def _string_to_version(cls, v: str | Version) -> Version: # noqa D102
|
||||
if isinstance(v, str):
|
||||
return Version(v)
|
||||
else:
|
||||
return v
|
||||
|
||||
|
||||
class ConfigMigrator:
|
||||
"""This class allows migrators to register their input and output versions."""
|
||||
@ -38,8 +29,8 @@ class ConfigMigrator:
|
||||
@classmethod
|
||||
def register(
|
||||
cls,
|
||||
from_version: Version | str,
|
||||
to_version: Version | str,
|
||||
from_version: str,
|
||||
to_version: str,
|
||||
) -> Callable[[Callable[[AppConfigDict], AppConfigDict]], Callable[[AppConfigDict], AppConfigDict]]:
|
||||
"""Define a decorator which registers the migration between two versions."""
|
||||
|
||||
@ -48,7 +39,9 @@ class ConfigMigrator:
|
||||
raise ValueError(
|
||||
f"function {function.__name__} is trying to register a migration for version {str(from_version)}, but this migration has already been registered."
|
||||
)
|
||||
cls._migrations.append(MigrationEntry(from_version=from_version, to_version=to_version, function=function))
|
||||
cls._migrations.append(
|
||||
MigrationEntry(from_version=Version(from_version), to_version=Version(to_version), function=function)
|
||||
)
|
||||
return function
|
||||
|
||||
return decorator
|
||||
|
Loading…
Reference in New Issue
Block a user