Coverage for instanovo/utils/file_downloader.py: 40%

1import os 

2from pathlib import Path 

3 

4import requests 

5from tqdm import tqdm 

6 

7from instanovo.__init__ import console 

8from instanovo.utils.colorlogging import ColorLog 

9 

10logger = ColorLog(console, __name__).logger 

11 

12 

13def download_file(url: str, path: Path, model_id: str, file_name: str) -> None: 

14 """Download a file with a progress bar. 

15 

16 Args: 

17 url (str): 

18 The URL to download the file from. 

19 path (Path): 

20 The path to save the file to. 

21 model_id (str): 

22 The model ID. 

23 file_name (str): 

24 The name of the file to download. 

25 """ 

26 # If not cached, download the file with a progress bar 

27 response = requests.get(url, stream=True) 

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

29 logger.info(f"Downloading model {model_id} from {url}") 

30 

31 if response.status_code == 404: 

32 raise ValueError(f"Model {model_id} not found at {url}. Status code: {response.status_code}") 

33 

34 with ( 

35 open(path, "wb") as file, 

36 tqdm( 

37 desc=file_name, 

38 total=total_size, 

39 unit="iB", 

40 unit_scale=True, 

41 unit_divisor=1024, 

42 ) as progress_bar, 

43 ): 

44 for data in response.iter_content(chunk_size=1024): 

45 size = file.write(data) 

46 progress_bar.update(size) 

47 if not os.path.getsize(path) == total_size: 

48 raise ValueError( 

49 f"Downloaded file is incomplete. Expected size of {total_size} bytes does not match downloaded size of {os.path.getsize(path)} bytes." 

50 ) 

51 

52 logger.info(f"Cached model {model_id} at {path}")