Coverage for src/edwh_files_plugin/files_plugin.py: 0%

111 statements  

« prev     ^ index     » next       coverage.py v7.6.10, created at 2025-02-28 11:40 +0100

1import json 

2import sys 

3import tempfile 

4import typing 

5from pathlib import Path 

6from typing import Optional 

7 

8import requests 

9from edwh import improved_task as task 

10from invoke import Context 

11 

12# rich.progress is fancier but much slower (100ms import) 

13# so use simpler progress library (also used by pip, before rich): 

14from progress.bar import ChargingBar 

15from requests_toolbelt.multipart.encoder import MultipartEncoder, MultipartEncoderMonitor 

16from rich import print # noqa: A004 

17from threadful import thread 

18from threadful.bonus import animate 

19 

20from edwh_files_plugin.compression import DEFAULT_COMPRESSION_LEVEL, Compression 

21 

22DEFAULT_TRANSFERSH_SERVER = "https://files.edwh.nl" 

23 

24 

25def require_protocol(url: str) -> str: 

26 """ 

27 Make sure 'url' has an HTTP or HTTPS schema. 

28 """ 

29 return url if url.startswith(("http://", "https://")) else f"https://{url}" 

30 

31 

32def create_callback(encoder: MultipartEncoder) -> typing.Callable[[MultipartEncoderMonitor], None]: 

33 """ 

34 Creates a callback function for monitoring the progress of an upload. 

35 

36 Args: 

37 encoder (MultipartEncoder): The multipart encoder that is uploading the data. 

38 

39 Returns: 

40 A callback function that updates a progress bar based on the amount of data that has been read by the encoder. 

41 

42 Example: 

43 

44 import requests 

45 from requests_toolbelt.multipart.encoder import MultipartEncoder, MultipartEncoderMonitor 

46 

47 def my_callback(monitor: MultipartEncoder): 

48 print('bytes sent: {0}'.format(monitor.bytes_read)) 

49 

50 filename = 'my_file.txt' 

51 

52 with open(filename, 'rb') as f: 

53 encoder = MultipartEncoder( 

54 fields={ 

55 'file': (filename, f, "text/plain"), 

56 } 

57 ) 

58 

59 monitor = MultipartEncoderMonitor(encoder, my_callback) 

60 result = requests.post('http://some-url.com/upload', 

61 data=monitor, 

62 headers={'Content-Type': monitor.content_type}) 

63 

64 """ 

65 bar = ChargingBar("Uploading", max=encoder.len) 

66 

67 def callback(monitor: MultipartEncoderMonitor) -> None: 

68 # goto instead of next because chunk size is unknown 

69 bar.goto(monitor.bytes_read) 

70 

71 return callback 

72 

73 

74# auto = best for directory; none for files 

75# gzip: .tgz for directory; .gz for files. Pigz or Gzip based on availability 

76FullCompressionTypes: typing.TypeAlias = typing.Literal["auto", "gzip", "zip", "tgz", "gz", "none"] 

77CliCompressionTypes: typing.TypeAlias = typing.Literal["auto", "gzip", "zip", "none"] 

78 

79 

80def upload_file( 

81 url: str, 

82 filename: str, 

83 filepath: Path, 

84 headers: Optional[dict[str, typing.Any]] = None, 

85 compression: FullCompressionTypes = "auto", 

86 compression_level: int = DEFAULT_COMPRESSION_LEVEL, 

87) -> requests.Response: 

88 """ 

89 Upload a file to an url. 

90 """ 

91 if headers is None: 

92 headers = {} 

93 

94 with tempfile.TemporaryDirectory() as tmpdir: 

95 if compression != "auto": 

96 new_filepath = Path(tmpdir) / filepath.name 

97 filename = compress_directory( 

98 filepath, 

99 new_filepath, 

100 extension="gz" if compression == "gzip" else compression, 

101 compression_level=compression_level, 

102 ) 

103 filepath = new_filepath 

104 

105 with filepath.open("rb") as f: 

106 encoder = MultipartEncoder( 

107 fields={ 

108 filename: (filename, f, "text/plain"), 

109 } 

110 ) 

111 

112 monitor = MultipartEncoderMonitor(encoder, create_callback(encoder)) 

113 return requests.post(url, data=monitor, headers=headers | {"Content-Type": monitor.content_type}) 

114 

115 

116@thread() 

117def _compress_directory( 

118 dir_path: str | Path, 

119 file_path: str | Path, 

120 extension: FullCompressionTypes = "auto", 

121 compression_level: int = DEFAULT_COMPRESSION_LEVEL, 

122) -> str: 

123 """ 

124 Compress a directory into a compressed (zip, gz) file. 

125 """ 

126 compressor = Compression.best() if extension == "auto" else Compression.for_extension(extension) 

127 

128 if not compressor: 

129 print(f"[red] No compression available for {extension} [/red]") 

130 print(f"[blue] Please choose one of : {Compression.available()}[/blue]") 

131 raise RuntimeError("Something went wrong during compression!") 

132 

133 if compressor.compress(dir_path, file_path, level=compression_level): 

134 return compressor.filename(dir_path) 

135 else: 

136 raise RuntimeError("Something went wrong during compression!") 

137 

138 

139def compress_directory( 

140 dir_path: str | Path, 

141 file_path: str | Path, 

142 extension: FullCompressionTypes = "auto", 

143 compression_level: int = DEFAULT_COMPRESSION_LEVEL, 

144) -> str: 

145 """ 

146 Compress a directory into a compressed file (zip, gz) and show a spinning animation. 

147 """ 

148 return animate( 

149 _compress_directory(dir_path, file_path, extension, compression_level=compression_level), 

150 text=f"Compressing directory {dir_path}", 

151 ) 

152 

153 

154def upload_directory( 

155 url: str, 

156 filepath: Path, 

157 headers: Optional[dict[str, typing.Any]] = None, 

158 upload_filename: Optional[str] = None, 

159 compression: FullCompressionTypes = "auto", 

160 compression_level: int = DEFAULT_COMPRESSION_LEVEL, 

161) -> requests.Response: 

162 """ 

163 Zip a directory and upload it to an url. 

164 

165 Args: 

166 url: which transfer.sh server to use 

167 filepath: which directory to upload 

168 headers: upload options 

169 upload_filename: by default, the directory name with compression extension (e.g. .gz, .zip) will be used 

170 compression: which method for compression to use (or best available by default) 

171 compression_level: The compression level is a measure of the compression quality (file size; int 0 - 9). 

172 """ 

173 filepath = filepath.expanduser().absolute() 

174 filename = filepath.resolve().name 

175 

176 with tempfile.TemporaryDirectory() as tmpdir: 

177 archive_path = Path(tmpdir) / filename 

178 compressed_filename = compress_directory( 

179 filepath, 

180 archive_path, 

181 extension="tgz" if compression == "gzip" else compression, 

182 compression_level=compression_level, 

183 ) # -> filename.zip e.g. 

184 

185 upload_filename = upload_filename or compressed_filename 

186 

187 return upload_file( 

188 url, upload_filename, Path(archive_path), headers=headers, compression_level=compression_level 

189 ) 

190 

191 

192@task(aliases=("add", "send")) 

193def upload( 

194 _: Context, 

195 filename: str | Path, 

196 server: str = DEFAULT_TRANSFERSH_SERVER, 

197 max_downloads: Optional[int] = None, 

198 max_days: Optional[int] = None, 

199 encrypt: Optional[str] = None, 

200 rename: Optional[str] = None, 

201 compression: CliCompressionTypes = "auto", # auto | pigz | gzip | zip 

202 compression_level: int = DEFAULT_COMPRESSION_LEVEL, 

203) -> None: 

204 """ 

205 Upload a file. 

206 

207 Args: 

208 _: invoke Context 

209 filename (str): path to the file to upload 

210 server (str): which transfer.sh server to use 

211 max_downloads (int): how often can the file be downloaded? 

212 max_days (int): how many days can the file be downloaded? 

213 encrypt (str): encryption password 

214 rename (str): upload the file/folder with a different name than it currently has 

215 compression (str): by default files are not compressed. 

216 For folders it will try pigz (.tgz), gzip (.tgz) then .zip. 

217 You can also explicitly specify a compression method for files and directory, 

218 and nothing else will be tried. 

219 compression_level (int): The compression level is a measure of the compression quality (file size). 

220 It is expressed as an integer in the range 1 - 9. 

221 Compression quality and performance are conflicting goals. 

222 Compression level 1 provides best performance at the expense of quality. 

223 Compression level 9 provides the smallest file size. 

224 """ 

225 headers: dict[str, str | int] = {} 

226 

227 if max_downloads: 

228 headers["Max-Downloads"] = max_downloads 

229 if max_days: 

230 headers["Max-Days"] = max_days 

231 if encrypt: 

232 headers["X-Encrypt-Password"] = encrypt 

233 

234 url = require_protocol(server) 

235 

236 filepath = Path(filename) 

237 

238 try: 

239 if filepath.is_dir(): 

240 response = upload_directory( 

241 url, 

242 filepath, 

243 headers, 

244 upload_filename=rename, 

245 compression=compression, 

246 compression_level=compression_level, 

247 ) 

248 else: 

249 response = upload_file( 

250 url, 

251 rename or str(filename), 

252 filepath, 

253 headers, 

254 compression=compression, 

255 compression_level=compression_level, 

256 ) 

257 except RuntimeError as e: 

258 print(f"[red] {e} [/red]") 

259 exit(1) 

260 

261 download_url = response.text.strip() 

262 delete_url = response.headers.get("x-url-delete") 

263 

264 print( 

265 json.dumps( 

266 { 

267 "status": response.status_code, 

268 "url": download_url, 

269 "delete": delete_url, 

270 "download_command": f"edwh file.download {download_url}", 

271 "delete_command": f"edwh file.delete {delete_url}", 

272 }, 

273 indent=2, 

274 ), 

275 ) 

276 

277 

278@task(aliases=("get", "receive")) 

279def download( 

280 _: Context, 

281 download_url: str, 

282 output_file: Optional[str | Path] = None, 

283 decrypt: Optional[str] = None, 

284 unpack: bool = False, 

285) -> None: 

286 """ 

287 Download a file. 

288 

289 Args: 

290 _ (Context) 

291 download_url (str): file to download 

292 output_file (str): path to store the file in 

293 decrypt (str): decryption token 

294 unpack (bool): unpack archive to file(s), removing the archive afterwards 

295 """ 

296 if output_file is None: 

297 output_file = download_url.split("/")[-1] 

298 output_path = Path(output_file) 

299 

300 download_url = require_protocol(download_url) 

301 

302 headers = {} 

303 if decrypt: 

304 headers["X-Decrypt-Password"] = decrypt 

305 

306 response = requests.get(download_url, headers=headers, stream=True) 

307 

308 if response.status_code >= 400: 

309 print("[red] Something went wrong: [/red]", response.status_code, response.content.decode(), file=sys.stderr) 

310 return 

311 

312 total = int(response.headers["Content-Length"]) // 1024 

313 with ( 

314 output_path.open("wb") as f, 

315 ): # <- open file when we're sure the status code is successful! 

316 for chunk in ChargingBar("Downloading", max=total).iter(response.iter_content(chunk_size=1024)): 

317 f.write(chunk) 

318 

319 if unpack: 

320 do_unpack(_, str(output_path), remove=True) 

321 

322 

323@task(aliases=("remove",)) 

324def delete(_: Context, deletion_url: str) -> None: 

325 """ 

326 Delete an uploaded file. 

327 

328 Args: 

329 _ (Context) 

330 deletion_url (str): File url + deletion token (from `x-url-delete`, shown in file.upload output) 

331 """ 

332 deletion_url = require_protocol(deletion_url) 

333 

334 response = requests.delete(deletion_url, timeout=15) 

335 

336 print( 

337 { 

338 "status": response.status_code, 

339 "response": response.text.strip(), 

340 } 

341 ) 

342 

343 

344@task(name="unpack") 

345def do_unpack(_: Context, filename: str, remove: bool = False) -> None: 

346 """ 

347 Decompress a given file. 

348 

349 Args: 

350 _: invoke Context 

351 filename (str): Path of the file to be decompressed 

352 remove (bool, optional): If True, original compressed file will be deleted after decompression. 

353 Defaults to False. 

354 

355 Returns: 

356 None 

357 """ 

358 filepath = Path(filename) 

359 ext = filepath.suffix 

360 

361 compressor = Compression.for_extension(ext) 

362 

363 if compressor and compressor.decompress(filepath, filepath.with_suffix("")): 

364 if remove: 

365 filepath.unlink() 

366 else: 

367 print("[red] Something went wrong unpacking! [/red]")