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
« 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
8import requests
9from edwh import improved_task as task
10from invoke import Context
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
20from edwh_files_plugin.compression import DEFAULT_COMPRESSION_LEVEL, Compression
22DEFAULT_TRANSFERSH_SERVER = "https://files.edwh.nl"
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}"
32def create_callback(encoder: MultipartEncoder) -> typing.Callable[[MultipartEncoderMonitor], None]:
33 """
34 Creates a callback function for monitoring the progress of an upload.
36 Args:
37 encoder (MultipartEncoder): The multipart encoder that is uploading the data.
39 Returns:
40 A callback function that updates a progress bar based on the amount of data that has been read by the encoder.
42 Example:
44 import requests
45 from requests_toolbelt.multipart.encoder import MultipartEncoder, MultipartEncoderMonitor
47 def my_callback(monitor: MultipartEncoder):
48 print('bytes sent: {0}'.format(monitor.bytes_read))
50 filename = 'my_file.txt'
52 with open(filename, 'rb') as f:
53 encoder = MultipartEncoder(
54 fields={
55 'file': (filename, f, "text/plain"),
56 }
57 )
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})
64 """
65 bar = ChargingBar("Uploading", max=encoder.len)
67 def callback(monitor: MultipartEncoderMonitor) -> None:
68 # goto instead of next because chunk size is unknown
69 bar.goto(monitor.bytes_read)
71 return callback
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"]
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 = {}
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
105 with filepath.open("rb") as f:
106 encoder = MultipartEncoder(
107 fields={
108 filename: (filename, f, "text/plain"),
109 }
110 )
112 monitor = MultipartEncoderMonitor(encoder, create_callback(encoder))
113 return requests.post(url, data=monitor, headers=headers | {"Content-Type": monitor.content_type})
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)
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!")
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!")
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 )
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.
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
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.
185 upload_filename = upload_filename or compressed_filename
187 return upload_file(
188 url, upload_filename, Path(archive_path), headers=headers, compression_level=compression_level
189 )
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.
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] = {}
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
234 url = require_protocol(server)
236 filepath = Path(filename)
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)
261 download_url = response.text.strip()
262 delete_url = response.headers.get("x-url-delete")
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 )
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.
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)
300 download_url = require_protocol(download_url)
302 headers = {}
303 if decrypt:
304 headers["X-Decrypt-Password"] = decrypt
306 response = requests.get(download_url, headers=headers, stream=True)
308 if response.status_code >= 400:
309 print("[red] Something went wrong: [/red]", response.status_code, response.content.decode(), file=sys.stderr)
310 return
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)
319 if unpack:
320 do_unpack(_, str(output_path), remove=True)
323@task(aliases=("remove",))
324def delete(_: Context, deletion_url: str) -> None:
325 """
326 Delete an uploaded file.
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)
334 response = requests.delete(deletion_url, timeout=15)
336 print(
337 {
338 "status": response.status_code,
339 "response": response.text.strip(),
340 }
341 )
344@task(name="unpack")
345def do_unpack(_: Context, filename: str, remove: bool = False) -> None:
346 """
347 Decompress a given file.
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.
355 Returns:
356 None
357 """
358 filepath = Path(filename)
359 ext = filepath.suffix
361 compressor = Compression.for_extension(ext)
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]")