Coverage for sparkle/CLI/help/snapshot_help.py: 84%

58 statements  

« prev     ^ index     » next       coverage.py v7.6.10, created at 2025-01-07 15:22 +0000

1#!/usr/bin/env python3 

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

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

4import shutil 

5import sys 

6import os 

7import time 

8from pathlib import Path 

9import zipfile 

10 

11from sparkle.CLI.help import global_variables as gv 

12from sparkle.tools.general import get_time_pid_random_string 

13 

14 

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

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

17 if name is None: 

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

19 try: 

20 login = os.getlogin() 

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

22 login = "unknown" 

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

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

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

26 available_dirs = [p.name for p in Path.cwd().iterdir()] 

27 root_working_dirs = [p for p in gv.settings().DEFAULT_working_dirs 

28 if p.name in available_dirs] 

29 for working_dir in root_working_dirs: 

30 if working_dir.exists(): 

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

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

33 shutil.rmtree(snapshot_tmp_path) 

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

35 

36 

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

38 """Remove the current Sparkle platform.""" 

39 filter = [] if filter is None else filter 

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

41 if working_dir not in filter: 

42 shutil.rmtree(working_dir, ignore_errors=True) 

43 

44 

45def create_working_dirs() -> None: 

46 """Create working directories.""" 

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

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

49 

50 

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

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

53 

54 Args: 

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

56 """ 

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

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

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

60 zip_ref.extractall(tmp_directory) 

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

62 shutil.rmtree(tmp_directory) 

63 

64 

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

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

67 

68 Args: 

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

70 """ 

71 if not snapshot_file.exists(): 

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

73 sys.exit(-1) 

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

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

76 sys.exit(-1) 

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

78 remove_current_platform() 

79 print("Existing Sparkle platform cleaned!") 

80 

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

82 extract_snapshot(snapshot_file) 

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

84 missing_dirs = [wd.name for wd in gv.settings().DEFAULT_working_dirs 

85 if not wd.exists()] 

86 print("ERROR: Failed to load Sparkle platform! The snapshot file may be outdated" 

87 " or corrupted. Missing the following directories: " 

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

89 sys.exit(-1) 

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