Coverage for src/sparkle/CLI/help/snapshot_help.py: 85%

1#!/usr/bin/env python3 

2# -*- coding: UTF-8 -*- 

3"""Helper functions to record and restore a Sparkle platform.""" 

4 

5import shutil 

6import sys 

7import os 

8import time 

9from pathlib import Path 

10import zipfile 

11 

12from sparkle.CLI.help import global_variables as gv 

13from sparkle.tools.general import get_time_pid_random_string 

14from sparkle.platform import Settings 

15 

16 

17def save_current_platform(name: str = None) -> None: 

18 """Store the current Sparkle platform in a .zip file.""" 

19 if name is None: 

20 time_stamp = time.strftime("%Y-%m-%d-%H.%M.%S", time.localtime(time.time())) 

21 try: 

22 login = os.getlogin() 

23 except Exception: # Can fail on for example CI pipelines 

24 login = "unknown" 

25 name = f"Snapshot_{login}_{time_stamp}" 

26 snapshot_tmp_path = gv.settings().DEFAULT_snapshot_dir / name 

27 snapshot_tmp_path.mkdir(parents=True) # Create temporary directory for zip 

28 available_dirs = [path.name for path in Path.cwd().iterdir()] 

29 root_working_dirs = [ 

30 path 

31 for path in gv.settings().DEFAULT_working_dirs 

32 if path.name in available_dirs 

33 ] 

34 for working_dir in root_working_dirs: 

35 if working_dir.exists(): 

36 shutil.copytree(working_dir, snapshot_tmp_path / working_dir.name) 

37 shutil.make_archive(snapshot_tmp_path, "zip", snapshot_tmp_path) 

38 shutil.rmtree(snapshot_tmp_path) 

39 print(f"Snapshot file {snapshot_tmp_path}.zip saved successfully!") 

40 

41 

42def remove_current_platform(filter: list[Path] = None) -> None: 

43 """Remove the current Sparkle platform.""" 

44 filter = [] if filter is None else filter 

45 for working_dir in gv.settings().DEFAULT_working_dirs: 

46 if working_dir not in filter: 

47 shutil.rmtree(working_dir, ignore_errors=True) 

48 Settings.DEFAULT_previous_settings_path.unlink(missing_ok=True) 

49 

50 

51def create_working_dirs() -> None: 

52 """Create working directories.""" 

53 for working_dir in gv.settings().DEFAULT_working_dirs: 

54 working_dir.mkdir(parents=True, exist_ok=True) 

55 

56 

57def extract_snapshot(snapshot_file: Path) -> None: 

58 """Restore a Sparkle platform from a snapshot. 

59 

60 Args: 

61 snapshot_file: Path to the where the current Sparkle platform should be stored. 

62 """ 

63 tmp_directory = Path(f"tmp_directory_{get_time_pid_random_string()}") 

64 gv.settings().DEFAULT_tmp_output.mkdir(exist_ok=True) 

65 with zipfile.ZipFile(snapshot_file, "r") as zip_ref: 

66 zip_ref.extractall(tmp_directory) 

67 shutil.copytree(tmp_directory, "./", dirs_exist_ok=True) 

68 shutil.rmtree(tmp_directory) 

69 

70 

71def load_snapshot(snapshot_file: Path) -> None: 

72 """Load a Sparkle platform from a snapshot. 

73 

74 Args: 

75 snapshot_file: File path to the file where the Sparkle platform is stored. 

76 """ 

77 if not snapshot_file.exists(): 

78 print(f"ERROR: Snapshot file {snapshot_file} does not exist!") 

79 sys.exit(-1) 

80 if not snapshot_file.suffix == ".zip": 

81 print(f"ERROR: File {snapshot_file} is not a .zip file!") 

82 sys.exit(-1) 

83 print("Cleaning existing Sparkle platform ...") 

84 remove_current_platform() 

85 print("Existing Sparkle platform cleaned!") 

86 

87 print(f"Loading snapshot file {snapshot_file} ...") 

88 extract_snapshot(snapshot_file) 

89 if any([not wd.exists() for wd in gv.settings().DEFAULT_working_dirs]): 

90 missing_dirs = [ 

91 wd.name for wd in gv.settings().DEFAULT_working_dirs if not wd.exists() 

92 ] 

93 print( 

94 "ERROR: Failed to load Sparkle platform! The snapshot file may be outdated" 

95 " or corrupted. Missing the following directories: " 

96 f"{', '.join(missing_dirs)}" 

97 ) 

98 sys.exit(-1) 

99 print(f"Snapshot file {snapshot_file} loaded successfully!")