diff --git a/antarest/study/storage/study_upgrader/__init__.py b/antarest/study/storage/study_upgrader/__init__.py index 242f84f519..6093f828e5 100644 --- a/antarest/study/storage/study_upgrader/__init__.py +++ b/antarest/study/storage/study_upgrader/__init__.py @@ -32,34 +32,18 @@ class UpgradeMethod(NamedTuple): files: List[Path] +# fmt: off UPGRADE_METHODS = [ - UpgradeMethod( - "700", "710", upgrade_710, [Path("settings/generaldata.ini")] - ), + UpgradeMethod("700", "710", upgrade_710, [Path("settings/generaldata.ini")]), UpgradeMethod("710", "720", upgrade_720, []), - UpgradeMethod( - "720", "800", upgrade_800, [Path("settings/generaldata.ini")] - ), - UpgradeMethod( - "800", - "810", - upgrade_810, - [Path("settings/generaldata.ini"), Path("input")], - ), + UpgradeMethod("720", "800", upgrade_800, [Path("settings/generaldata.ini")]), + UpgradeMethod("800", "810", upgrade_810, [Path("settings/generaldata.ini"), Path("input")]), UpgradeMethod("810", "820", upgrade_820, [Path("input/links")]), - UpgradeMethod( - "820", - "830", - upgrade_830, - [Path("settings/generaldata.ini"), Path("input/areas")], - ), - UpgradeMethod( - "830", "840", upgrade_840, [Path("settings/generaldata.ini")] - ), - UpgradeMethod( - "840", "850", upgrade_850, [Path("settings/generaldata.ini")] - ), + UpgradeMethod("820", "830", upgrade_830, [Path("settings/generaldata.ini"), Path("input/areas")]), + UpgradeMethod("830", "840", upgrade_840, [Path("settings/generaldata.ini")]), + UpgradeMethod("840", "850", upgrade_850, [Path("settings/generaldata.ini")]), ] +# fmt: on class InvalidUpgrade(HTTPException): @@ -93,24 +77,9 @@ def upgrade_study(study_path: Path, target_version: str) -> None: try: src_version = get_current_version(study_path) files_to_upgrade = can_upgrade_version(src_version, target_version) - files_to_upgrade.append(Path("study.antares")) - files_to_retrieve = [] - for path in files_to_upgrade: - entire_path = study_path / path - if entire_path.is_dir(): - if not (tmp_dir / path).exists(): - shutil.copytree( - entire_path, tmp_dir / path, dirs_exist_ok=True - ) - files_to_retrieve.append(path) - elif len(path.parts) == 1: - shutil.copy(entire_path, tmp_dir / path) - files_to_retrieve.append(path) - else: - parent_path = path.parent - (tmp_dir / parent_path).mkdir(parents=True) - shutil.copy(entire_path, tmp_dir / parent_path) - files_to_retrieve.append(path) + files_to_retrieve = _copies_only_necessary_files( + files_to_upgrade, study_path, tmp_dir + ) _do_upgrade(tmp_dir, src_version, target_version) except (StudyValidationError, InvalidUpgrade) as e: shutil.rmtree(tmp_dir) @@ -121,19 +90,7 @@ def upgrade_study(study_path: Path, target_version: str) -> None: logger.error(f"Unhandled exception : {e}", exc_info=True) raise else: - for k, path in enumerate(files_to_retrieve): - backup_dir = Path( - tempfile.mkdtemp( - suffix=f".backup_{k}.tmp", - prefix="~", - dir=study_path.parent, - ) - ) - backup_dir.rmdir() - original_path = study_path / path - original_path.rename(backup_dir) - (tmp_dir / path).rename(original_path) - shutil.rmtree(backup_dir, ignore_errors=True) + _replace_safely_original_files(files_to_retrieve, study_path, tmp_dir) def get_current_version(study_path: Path) -> str: @@ -171,6 +128,9 @@ def can_upgrade_version(from_version: str, to_version: str) -> List[Path]: from_version: The current version of the study. to_version: The target version of the study. + Returns: + If the upgrade is possible, the list of concerned folders and files + Raises: InvalidUpgrade: If the upgrade is not possible. """ @@ -229,6 +189,70 @@ def _update_study_antares_file(target_version: str, study_path: Path) -> None: file.write_text(content, encoding="utf-8") +def _copies_only_necessary_files( + files_to_upgrade: List[Path], study_path: Path, tmp_path: Path +) -> List[Path]: + """ + Copies files concerned by the version upgrader into a temporary directory. + + Args: + study_path: Path to the study. + tmp_path: Path to the temporary directory where the file modification will be performed. + files_to_upgrade: List[Path]: List of the files and folders concerned by the upgrade. + + Returns: + The list of files and folders that were really copied. It's the same as files_to_upgrade but + without any children that has parents already in the list. + """ + files_to_upgrade.append(Path("study.antares")) + files_to_retrieve = [] + for path in files_to_upgrade: + entire_path = study_path / path + if entire_path.is_dir(): + if not (tmp_path / path).exists(): + shutil.copytree( + entire_path, tmp_path / path, dirs_exist_ok=True + ) + files_to_retrieve.append(path) + elif len(path.parts) == 1: + shutil.copy(entire_path, tmp_path / path) + files_to_retrieve.append(path) + else: + parent_path = path.parent + (tmp_path / parent_path).mkdir(parents=True) + shutil.copy(entire_path, tmp_path / parent_path) + files_to_retrieve.append(path) + return files_to_retrieve + + +def _replace_safely_original_files( + files_to_replace: List[Path], study_path: Path, tmp_path: Path +) -> None: + """ + Replace files/folders of the study that should be upgraded by their copy already upgraded in the tmp directory. + It uses Path.rename() and an intermediary tmp directory to swap the folders safely. + In the end, all tmp directories are removed. + + Args: + study_path: Path to the study. + tmp_path: Path to the temporary directory where the file modification will be performed. + files_to_replace: List[Path]: List of files and folders that were really copied, cf. _copies_only_necessary_files's doc (just above). + """ + for k, path in enumerate(files_to_replace): + backup_dir = Path( + tempfile.mkdtemp( + suffix=f".backup_{k}.tmp", + prefix="~", + dir=study_path.parent, + ) + ) + backup_dir.rmdir() + original_path = study_path / path + original_path.rename(backup_dir) + (tmp_path / path).rename(original_path) + shutil.rmtree(backup_dir, ignore_errors=True) + + def _do_upgrade( study_path: Path, src_version: str, target_version: str ) -> None: