Coverage for hopwise/utils/url.py: 23%

1"""hopwise.utils.url 

2################################ 

3Reference code: 

4 https://github.com/snap-stanford/ogb/blob/master/ogb/utils/url.py 

5""" 

6 

7import errno 

8import os 

9import os.path as osp 

10import urllib.request as ur 

11import zipfile 

12from logging import getLogger 

13 

14from hopwise.utils import progress_bar 

15 

16GBFACTOR = float(1 << 30) 

17 

18 

19def decide_download(url): 

20 d = ur.urlopen(url) 

21 size = int(d.info()["Content-Length"]) / GBFACTOR 

22 

23 ### confirm if larger than 1GB 

24 if size > 1: 

25 return input("This will download %.2fGB. Will you proceed? (y/N)\n" % (size)).lower() == "y" 

26 else: 

27 return True 

28 

29 

30def makedirs(path): 

31 try: 

32 os.makedirs(osp.expanduser(osp.normpath(path))) 

33 except OSError as e: 

34 if e.errno != errno.EEXIST and osp.isdir(path): 

35 raise e 

36 

37 

38def download_url(url, folder): 

39 """Downloads the content of an URL to a specific folder. 

40 

41 Args: 

42 url (string): The url. 

43 folder (string): The folder. 

44 """ 

45 filename = url.rpartition("/")[2] 

46 path = osp.join(folder, filename) 

47 logger = getLogger() 

48 

49 if osp.exists(path) and osp.getsize(path) > 0: # pragma: no cover 

50 logger.info(f"Using exist file {filename}") 

51 return path 

52 

53 logger.info(f"Downloading {url}") 

54 

55 makedirs(folder) 

56 data = ur.urlopen(url) 

57 

58 size = int(data.info()["Content-Length"]) 

59 

60 chunk_size = 1024 * 1024 

61 num_iter = int(size / chunk_size) + 2 

62 

63 downloaded_size = 0 

64 

65 try: 

66 with open(path, "wb") as f: 

67 pbar = progress_bar(range(num_iter)) 

68 for i in pbar: 

69 chunk = data.read(chunk_size) 

70 downloaded_size += len(chunk) 

71 pbar.set_description(f"Downloaded {float(downloaded_size) / GBFACTOR:.2f} GB") 

72 f.write(chunk) 

73 except Exception: 

74 if os.path.exists(path): 

75 os.remove(path) 

76 raise RuntimeError("Stopped downloading due to interruption.") 

77 

78 return path 

79 

80 

81def extract_zip(path, folder): 

82 """Extracts a zip archive to a specific folder. 

83 

84 Args: 

85 path (string): The path to the tar archive. 

86 folder (string): The folder. 

87 """ 

88 logger = getLogger() 

89 logger.info(f"Extracting {path}") 

90 with zipfile.ZipFile(path, "r") as f: 

91 f.extractall(folder) 

92 

93 

94def rename_atomic_files(folder, old_name, new_name): 

95 """Rename all atomic files in a given folder. 

96 

97 Args: 

98 folder (string): The folder. 

99 old_name (string): Old name for atomic files. 

100 new_name (string): New name for atomic files. 

101 """ 

102 files = os.listdir(folder) 

103 for f in files: 

104 base, suf = os.path.splitext(f) 

105 if old_name not in base: 

106 continue 

107 if suf not in {".inter", ".user", ".item"}: 

108 logger = getLogger() 

109 logger.warning(f"Moving downloaded file with suffix [{suf}].") 

110 os.rename( 

111 os.path.join(folder, f), 

112 os.path.join(folder, base.replace(old_name, new_name) + suf), 

113 ) 

114 

115 

116if __name__ == "__main__": 

117 pass