Skip to content

Commit

Permalink
use packaging.version rather than version-parse
Browse files Browse the repository at this point in the history
  • Loading branch information
Lincoln Stein committed Apr 19, 2024
1 parent 6ad1948 commit 36495b7
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 22 deletions.
32 changes: 11 additions & 21 deletions invokeai/app/services/config/config_migrate.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,32 +7,22 @@
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Callable, List, TypeVar

from packaging.version import Version
from pydantic import BaseModel, ConfigDict, field_validator
from version_parser import Version

if TYPE_CHECKING:
pass

AppConfigDict = TypeVar("AppConfigDict", bound=dict[str, Any])


class AppVersion(Version):
"""Stringlike object that sorts like a version."""

def __hash__(self) -> int: # noqa D105
return hash(str(self))

def __repr__(self) -> str: # noqa D105
return f"AppVersion('{str(self)}')"


class ConfigMigratorBase(ABC):
"""This class allows migrators to register their input and output versions."""

@classmethod
@abstractmethod
def register(
cls, from_version: AppVersion, to_version: AppVersion
cls, from_version: Version, to_version: Version
) -> Callable[[Callable[[AppConfigDict], AppConfigDict]], Callable[[AppConfigDict], AppConfigDict]]:
"""Define a decorator which registers the migration between two versions."""

Expand All @@ -54,15 +44,15 @@ class MigrationEntry(BaseModel):

model_config = ConfigDict(arbitrary_types_allowed=True)

from_version: AppVersion
to_version: AppVersion
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 | AppVersion) -> AppVersion: # noqa D102
def _string_to_version(cls, v: str | Version) -> Version: # noqa D102
if isinstance(v, str):
return AppVersion(v)
return Version(v)
else:
return v

Expand All @@ -75,8 +65,8 @@ class ConfigMigrator(ConfigMigratorBase):
@classmethod
def register(
cls,
from_version: AppVersion | str,
to_version: AppVersion | str,
from_version: Version | str,
to_version: Version | str,
) -> Callable[[Callable[[AppConfigDict], AppConfigDict]], Callable[[AppConfigDict], AppConfigDict]]:
"""Define a decorator which registers the migration between two versions."""

Expand All @@ -92,7 +82,7 @@ def decorator(function: Callable[[AppConfigDict], AppConfigDict]) -> Callable[[A

@staticmethod
def _check_for_overlaps(migrations: List[MigrationEntry]) -> None:
current_version = AppVersion("0.0.0")
current_version = Version("0.0.0")
for m in migrations:
if current_version > m.from_version:
raise ValueError(f"Version range overlap detected while processing function {m.function.__name__}")
Expand All @@ -116,9 +106,9 @@ def migrate(cls, config_dict: AppConfigDict) -> AppConfigDict:
cls._check_for_overlaps(sorted_migrations)

if "InvokeAI" in config_dict:
version = AppVersion("3.0.0")
version = Version("3.0.0")
else:
version = AppVersion(config_dict["schema_version"])
version = Version(config_dict["schema_version"])

for migration in sorted_migrations:
if version >= migration.from_version and version < migration.to_version:
Expand Down
1 change: 0 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,6 @@ dependencies = [
"semver~=3.0.1",
"send2trash",
"test-tube~=0.7.5",
"version-parser",
"windows-curses; sys_platform=='win32'",
]

Expand Down

0 comments on commit 36495b7

Please sign in to comment.