Coverage for kwave/__init__.py: 65%

1import hashlib 

2import json 

3import logging 

4import os 

5import platform 

6from pathlib import Path 

7from typing import List 

8from urllib.request import urlretrieve 

9 

10# Test installation with: 

11# python3 -m pip install -i https://test.pypi.org/simple/ --extra-index-url=https://pypi.org/simple/ k-Wave-python==0.3.0 

12__version__ = "0.4.1" 

13 

14# Constants and Configurations 

15URL_BASE = "https://github.com/waltsims/" 

16BINARY_VERSION = "v1.3.0" 

17PREFIX = f"{URL_BASE}kspaceFirstOrder-{ } -{ } /releases/download/{BINARY_VERSION}/" 

18PLATFORM = platform.system().lower() 

19 

20if PLATFORM not in ["linux", "windows", "darwin"]: 20 ↛ 21line 20 didn't jump to line 21 because the condition on line 20 was never true

21 raise NotImplementedError(f"k-wave-python is currently unsupported on this operating system: {PLATFORM}.") 

22 

23# TODO: install directly in to /bin/ directory system directory is no longer needed 

24# TODO: deprecate in 0.5.0 

25BINARY_PATH = Path(__file__).parent / "bin" / PLATFORM 

26BINARY_DIR = BINARY_PATH # add alias for BINARY_PATH for now 

27 

28 

29WINDOWS_DLLS = [ 

30 "cufft64_10.dll", 

31 "hdf5.dll", 

32 "hdf5_hl.dll", 

33 "libiomp5md.dll", 

34 "libmmd.dll", 

35 "msvcp140.dll", 

36 "svml_dispmd.dll", 

37 "szip.dll", 

38 "vcruntime140.dll", 

39 "zlib.dll", 

40] 

41 

42EXECUTABLE_PREFIX = "kspaceFirstOrder-" 

43ARCHITECTURES = ["omp", "cuda"] 

44 

45 

46def get_windows_release_urls(architecture: str) -> list: 

47 specific_filenames = [EXECUTABLE_PREFIX + architecture + ".exe"] 

48 if architecture == "omp": 

49 specific_filenames += WINDOWS_DLLS 

50 release_urls = [PREFIX.format(architecture.upper(), PLATFORM.lower()) + filename for filename in specific_filenames] 

51 return release_urls 

52 

53 

54URL_DICT = { 

55 "linux": { 

56 "cuda": [URL_BASE + f"kspaceFirstOrder-CUDA-{PLATFORM}/releases/download/v1.3.1/{EXECUTABLE_PREFIX}CUDA"], 

57 "omp": [URL_BASE + f"kspaceFirstOrder-OMP-{PLATFORM}/releases/download/{BINARY_VERSION}/{EXECUTABLE_PREFIX}OMP"], 

58 }, 

59 "darwin": { 

60 "cuda": [], 

61 "omp": [URL_BASE + f"k-wave-omp-{PLATFORM}/releases/download/v0.3.0rc2/{EXECUTABLE_PREFIX}OMP"], 

62 }, 

63 "windows": {architecture: get_windows_release_urls(architecture) for architecture in ARCHITECTURES}, 

64} 

65 

66 

67def _hash_file(filepath: str) -> str: 

68 buf_size = 65536 # 64kb chunks 

69 md5 = hashlib.md5() 

70 

71 with open(filepath, "rb") as f: 

72 while True: 

73 data = f.read(buf_size) 

74 if not data: 

75 break 

76 md5.update(data) 

77 return md5.hexdigest() 

78 

79 

80def _is_binary_present(binary_name: str, binary_type: str) -> bool: 

81 binary_filepath = BINARY_PATH / binary_name 

82 binary_file_exists = os.path.exists(binary_filepath) 

83 if not binary_file_exists: 83 ↛ 84line 83 didn't jump to line 84 because the condition on line 83 was never true

84 return False 

85 

86 if binary_type is None: 86 ↛ 89line 86 didn't jump to line 89 because the condition on line 86 was never true

87 # this is non-kwave windows binary 

88 # it already exists according to the check above 

89 return True 

90 existing_metadata_path = BINARY_PATH / f"{binary_name}_metadata.json" 

91 

92 if not os.path.exists(existing_metadata_path): 92 ↛ 96line 92 didn't jump to line 96 because the condition on line 92 was never true

93 # metadata does not exist => binaries may or may not exist 

94 # Let's play safe and claim they don't exist 

95 # This will trigger binary download and generation of binary metadata 

96 return False 

97 existing_metadata = json.loads(Path(existing_metadata_path).read_text()) 

98 

99 # If metadata was somehow corrupted 

100 file_hash = _hash_file(binary_filepath) 

101 if existing_metadata["file_hash"] != file_hash: 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true

102 return False 

103 

104 # If there is a new binary 

105 latest_urls = URL_DICT[PLATFORM][binary_type] 

106 if existing_metadata["url"] not in latest_urls: 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true

107 return False 

108 

109 # No need to check `version` field for now 

110 # because we version is already present in the URL 

111 return True 

112 

113 

114def binaries_present() -> bool: 

115 """ 

116 Check if binaries are present 

117 Returns: 

118 bool, True if binaries are present, False otherwise 

119 

120 """ 

121 binary_list = [] 

122 for binary_type in ARCHITECTURES: 

123 for binary_name in URL_DICT[PLATFORM][binary_type]: 

124 binary_list.append((binary_name.split("/")[-1], binary_type)) 

125 

126 missing_binaries: List[str] = [] 

127 

128 for binary_name, binary_type in binary_list: 

129 if not _is_binary_present(binary_name, binary_type): 129 ↛ 130line 129 didn't jump to line 130 because the condition on line 129 was never true

130 missing_binaries.append(binary_name) 

131 

132 if len(missing_binaries) > 0: 132 ↛ 133line 132 didn't jump to line 133 because the condition on line 132 was never true

133 missing_binaries_str = ", ".join(missing_binaries) 

134 logging.log( 

135 logging.INFO, 

136 f"Following binaries were not found: {missing_binaries_str}" 

137 "If this is first time you're running k-wave-python, " 

138 "binaries will be downloaded automatically.", 

139 ) 

140 

141 return len(missing_binaries) == 0 

142 

143 

144def _record_binary_metadata(binary_version: str, binary_filepath: str, binary_url: str, filename: str) -> None: 

145 # note: version is not immediately useful at the moment 

146 # because it is already present in the url and we use url to understand if versions match 

147 # However, let's record it anyway. Maybe it will be useful in the future. 

148 metadata = {"url": binary_url, "version": binary_version, "file_hash": _hash_file(binary_filepath)} 

149 metadata_filename = f"{filename}_metadata.json" 

150 metadata_filepath = BINARY_PATH / metadata_filename 

151 with open(metadata_filepath, "w") as outfile: 

152 json.dump(metadata, outfile, indent=4) 

153 

154 

155def download_binaries(system_os: str, bin_type: str): 

156 """ 

157 Download binary from release url 

158 Args: 

159 system_os: string, current system type 

160 bin_type: string of "OMP" or "CUDA" 

161 

162 Returns: 

163 None 

164 

165 """ 

166 for url in URL_DICT[system_os][bin_type]: 

167 # Extract the file name from the GitHub release URL 

168 binary_version, filename = url.split("/")[-2:] 

169 

170 logging.log(logging.INFO, f"Downloading {filename} to {BINARY_PATH}...") 

171 

172 # Create the directory if it does not yet exist 

173 os.makedirs(BINARY_PATH, exist_ok=True) 

174 

175 # Download the binary file 

176 try: 

177 binary_filepath = os.path.join(BINARY_PATH, filename) 

178 urlretrieve(url, binary_filepath) 

179 _record_binary_metadata(binary_version=binary_version, binary_filepath=binary_filepath, binary_url=url, filename=filename) 

180 

181 except TimeoutError: 

182 logging.log( 

183 logging.WARN, 

184 f"Download of {filename} timed out. " 

185 "This can be due to slow internet connection. " 

186 "Partially downloaded files will be removed.", 

187 ) 

188 try: 

189 os.remove(BINARY_PATH) 

190 except Exception: 

191 folder_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "bin") 

192 logging.warning( 

193 "Error occurred while removing partially downloaded binary. " 

194 f"Please manually delete the `{folder_path}` folder which " 

195 "can be found in your virtual environment." 

196 ) 

197 

198 

199def install_binaries(): 

200 for binary_type in ARCHITECTURES: 

201 download_binaries(PLATFORM, binary_type) 

202 

203 

204if not binaries_present(): 204 ↛ 205line 204 didn't jump to line 205 because the condition on line 204 was never true

205 install_binaries()