diff --git a/mlem/api/migrations.py b/mlem/api/migrations.py index bde896df..afd6020f 100644 --- a/mlem/api/migrations.py +++ b/mlem/api/migrations.py @@ -30,16 +30,21 @@ def migrate(path: str, project: Optional[str] = None, recursive: bool = False): _migrate_one(loc) -def _migrate_one(location: Location): - with location.open("r") as f: - payload = safe_load(f) - +def apply_migrations(payload: dict): changed = False for migration in _migrations: migrated = migration(payload) if migrated is not None: payload = migrated changed = True + return payload, changed + + +def _migrate_one(location: Location): + with location.open("r") as f: + payload = safe_load(f) + + payload, changed = apply_migrations(payload) if changed: echo(f"Migrated MLEM Object at {location}") diff --git a/mlem/core/metadata.py b/mlem/core/metadata.py index 06fab034..465cd500 100644 --- a/mlem/core/metadata.py +++ b/mlem/core/metadata.py @@ -138,6 +138,7 @@ def load( rev: Optional[str] = None, batch_size: Optional[int] = None, follow_links: bool = True, + try_migrations: bool = False, ) -> Any: """Load python object saved by MLEM @@ -158,6 +159,7 @@ def load( rev=rev, follow_links=follow_links, load_value=batch_size is None, + try_migrations=try_migrations, ) if isinstance(meta, MlemData) and batch_size: return meta.read_batch(batch_size) @@ -175,6 +177,7 @@ def load_meta( follow_links: bool = True, load_value: bool = False, fs: Optional[AbstractFileSystem] = None, + try_migrations: bool = False, *, force_type: Literal[None] = None, ) -> MlemObject: @@ -189,6 +192,7 @@ def load_meta( follow_links: bool = True, load_value: bool = False, fs: Optional[AbstractFileSystem] = None, + try_migrations: bool = False, *, force_type: Optional[Type[T]] = None, ) -> T: @@ -203,6 +207,7 @@ def load_meta( follow_links: bool = True, load_value: bool = False, fs: Optional[AbstractFileSystem] = None, + try_migrations: bool = False, *, force_type: Optional[Type[T]] = None, ) -> T: @@ -216,6 +221,7 @@ def load_meta( actual object link points to. Defaults to True. load_value: Load actual python object incorporated in MlemObject. Defaults to False. fs: filesystem to load from. If not provided, will be inferred from path + try_migrations: If loading older versions of metadata, try to apply migrations force_type: type of meta to be loaded. Defaults to MlemObject (any mlem meta) Returns: MlemObject: Saved MlemObject @@ -232,6 +238,7 @@ def load_meta( meta = cls.read( location=find_meta_location(location), follow_links=follow_links, + try_migrations=try_migrations, ) log_meta_params(meta, add_object_type=True) if load_value: @@ -274,14 +281,21 @@ def find_meta_location(location: Location) -> Location: def list_objects( - path: str = ".", fs: Optional[AbstractFileSystem] = None, recursive=True + path: str = ".", + fs: Optional[AbstractFileSystem] = None, + recursive=True, + try_migrations=False, ) -> Dict[Type[MlemObject], List[MlemObject]]: loc = Location.resolve(path, fs=fs) result = defaultdict(list) postfix = f"/**{MLEM_EXT}" if recursive else f"/*{MLEM_EXT}" for filepath in loc.fs.glob(loc.fullpath + postfix, detail=False): meta = load_meta( - filepath, fs=loc.fs, load_value=False, follow_links=False + filepath, + fs=loc.fs, + load_value=False, + follow_links=False, + try_migrations=try_migrations, ) type_ = meta.__class__ if isinstance(meta, MlemLink): diff --git a/mlem/core/objects.py b/mlem/core/objects.py index 0cf03162..f69486d1 100644 --- a/mlem/core/objects.py +++ b/mlem/core/objects.py @@ -150,6 +150,7 @@ def read( cls: Type[T], location: Location, follow_links: bool = True, + try_migrations: bool = False, ) -> T: """ Read object in (path, fs) @@ -170,6 +171,10 @@ def read( ) with location.open() as f: payload = safe_load(f) + if try_migrations: + from mlem.api.migrations import apply_migrations + + payload, _ = apply_migrations(payload) res = parse_obj_as(MlemObject, payload).bind(location) if follow_links and isinstance(res, MlemLink): link = res.load_link() diff --git a/tests/api/test_migrations.py b/tests/api/test_migrations.py index 51e76117..54c7c0b1 100644 --- a/tests/api/test_migrations.py +++ b/tests/api/test_migrations.py @@ -61,3 +61,14 @@ def test_directory(tmpdir, old_data, new_data, recursive): assert load_meta(subdir_path) != new_data except ValidationError: pass + + +@pytest.mark.parametrize("old_data,new_data", [model_03]) +def test_load_with_migration(tmpdir, old_data, new_data): + path = tmpdir / "model.mlem" + path.write_text(safe_dump(old_data), encoding="utf8") + + meta = load_meta(path, try_migrations=True) + + assert isinstance(meta, MlemObject) + assert meta == new_data