Coverage for instanovo/scripts/get_zenodo_record.py: 96%

1# /// script 

2# requires-python = ">=3.10" 

3# dependencies = [ 

4# "requests", 

5# "typer", 

6# "rich", 

7# ] 

8# /// 

9 

10from __future__ import annotations 

11 

12import logging 

13import os 

14import zipfile 

15 

16import requests 

17import typer 

18from rich.console import Console 

19from rich.progress import ( 

20 BarColumn, 

21 DownloadColumn, 

22 Progress, 

23 TaskID, 

24 TimeRemainingColumn, 

25 TransferSpeedColumn, 

26) 

27 

28logger = logging.getLogger() 

29logger.setLevel(logging.INFO) 

30 

31 

32RECORD_ID = "17816199" 

33 

34 

35def get_zenodo(zenodo_url: str, zip_path: str, progress: Progress, task_id: TaskID) -> None: 

36 """Fetches specified zenodo record.""" 

37 try: 

38 response = requests.get(zenodo_url, stream=True) 

39 total_size = int(response.headers.get("content-length", 0)) 

40 progress.update(task_id, total=total_size) 

41 

42 with open(zip_path, "wb") as file: 

43 for data in response.iter_content(chunk_size=8192): 

44 file.write(data) 

45 progress.update(task_id, advance=len(data)) 

46 

47 logger.info(f"Zip file at url {zenodo_url} downloaded to {zip_path}.") 

48 

49 except requests.RequestException as e: 

50 logger.error(f"Failed to download the Zenodo record: {e}") 

51 raise 

52 

53 

54def unzip_zenodo(zip_path: str, extract_path: str) -> None: 

55 """Extracts zip file to specified location.""" 

56 try: 

57 os.makedirs(extract_path, exist_ok=True) 

58 

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

60 zip_ref.extractall(extract_path) 

61 

62 logger.info(f"Zip files extracted to {extract_path}") 

63 

64 except zipfile.BadZipFile as e: 

65 logger.error(f"Failed to unzip the file: {e}") 

66 raise 

67 

68 

69app = typer.Typer() 

70 

71 

72@app.command() 

73def main( 

74 zenodo_url: str = typer.Option( 

75 f"https://zenodo.org/records/{RECORD_ID}/files/instanovo_test_resources.zip", 

76 help="URL of the Zenodo record to download", 

77 ), 

78 zip_path: str = typer.Option( 

79 "./tests/instanovo_test_resources.zip", 

80 help="Path where the downloaded zip file will be saved", 

81 ), 

82 extract_path: str = typer.Option("./tests/instanovo_test_resources", help="Path where the zip file contents will be extracted"), 

83) -> None: 

84 """Downloads and extracts the zenodo record used for unit and integration tests.""" 

85 if os.path.exists(f"{extract_path}") and os.listdir(f"{extract_path}"): 

86 if os.path.exists(f"{extract_path}/record_id.txt"): 

87 with open(f"{extract_path}/record_id.txt", "r") as f: 

88 record_id = f.read().strip() 

89 if record_id == RECORD_ID: 

90 typer.echo(f"Record is up to date, skipping download and extraction. Path '{extract_path}' already exists and is non-empty.") 

91 raise typer.Exit() 

92 else: 

93 typer.echo("Record is outdated, downloading new record.") 

94 else: 

95 typer.echo("Record ID is not documented, downloading new record.") 

96 

97 console = Console() # Create a Console instance for Rich 

98 progress = Progress( 

99 DownloadColumn(), 

100 BarColumn(), 

101 "[progress.percentage]{task.percentage:>3.1f}%", 

102 "•", 

103 TransferSpeedColumn(), 

104 "•", 

105 TimeRemainingColumn(), 

106 console=console, # Pass the console instance to Progress 

107 ) 

108 

109 task_id = progress.add_task("download", filename=zip_path, start=False) 

110 

111 with progress: 

112 progress.start_task(task_id) 

113 get_zenodo(zenodo_url, zip_path, progress, task_id) 

114 

115 unzip_zenodo(zip_path, extract_path) 

116 

117 with open(f"{extract_path}/record_id.txt", "w") as f: 

118 f.write(RECORD_ID) 

119 

120 

121if __name__ == "__main__": 

122 app()