169 lines
4.4 KiB
Python
169 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import os
|
|
import os.path as osp
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
from collections.abc import Callable
|
|
from typing import TypedDict
|
|
|
|
if sys.version_info >= (3, 12):
|
|
from typing import Unpack
|
|
else:
|
|
from typing_extensions import Unpack
|
|
|
|
import filelock
|
|
|
|
from .download import download
|
|
|
|
|
|
class _DownloadKwargs(TypedDict, total=False):
|
|
proxy: str | None
|
|
speed: float | None
|
|
use_cookies: bool
|
|
verify: bool | str
|
|
id: str | None
|
|
resume: bool
|
|
format: str | None
|
|
user_agent: str | None
|
|
progress: Callable[[int, int | None], None] | None
|
|
|
|
|
|
cache_root = osp.join(osp.expanduser("~"), ".cache/gdown")
|
|
if not osp.exists(cache_root):
|
|
try:
|
|
os.makedirs(cache_root)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def cached_download(
|
|
url: str | None = None,
|
|
path: str | None = None,
|
|
quiet: bool = False,
|
|
postprocess: Callable[[str], object] | None = None,
|
|
hash: str | None = None,
|
|
**kwargs: Unpack[_DownloadKwargs],
|
|
) -> str:
|
|
"""Cached download from URL.
|
|
|
|
Parameters
|
|
----------
|
|
url:
|
|
URL. Google Drive URL is also supported.
|
|
path:
|
|
Output filename. Default is basename of URL.
|
|
quiet:
|
|
Suppress terminal output. Default is False.
|
|
postprocess:
|
|
Function called with filename as postprocess.
|
|
hash:
|
|
Hash value of file in the format of {algorithm}:{hash_value}
|
|
such as sha256:abcdef.... Supported algorithms: md5, sha1, sha256, sha512.
|
|
kwargs:
|
|
Keyword arguments to be passed to `download`.
|
|
|
|
Returns
|
|
-------
|
|
path:
|
|
Output filename.
|
|
|
|
Raises
|
|
------
|
|
ValueError
|
|
If url is not specified when path is not specified.
|
|
DownloadError
|
|
If the download fails.
|
|
"""
|
|
if path is None:
|
|
if url is None:
|
|
raise ValueError("url must be specified when path is not specified")
|
|
path = (
|
|
url.replace("/", "-SLASH-")
|
|
.replace(":", "-COLON-")
|
|
.replace("=", "-EQUAL-")
|
|
.replace("?", "-QUESTION-")
|
|
)
|
|
path = osp.join(cache_root, path)
|
|
|
|
# check existence
|
|
if osp.exists(path) and not hash:
|
|
if not quiet:
|
|
print(f"File exists: {path}", file=sys.stderr)
|
|
return path
|
|
elif osp.exists(path) and hash:
|
|
try:
|
|
_assert_filehash(path=path, hash=hash, quiet=quiet)
|
|
return path
|
|
except AssertionError as e:
|
|
print(e, file=sys.stderr)
|
|
|
|
# download
|
|
lock_path = osp.join(cache_root, "_dl_lock")
|
|
try:
|
|
os.makedirs(osp.dirname(path))
|
|
except OSError:
|
|
pass
|
|
temp_root = tempfile.mkdtemp(dir=cache_root)
|
|
try:
|
|
temp_path = osp.join(temp_root, "dl")
|
|
|
|
log_message_hash = f"Hash: {hash}\n" if hash else ""
|
|
download(
|
|
url=url,
|
|
output=temp_path,
|
|
quiet=quiet,
|
|
log_messages={
|
|
"start": f"Cached downloading...\n{log_message_hash}",
|
|
"output": f"To: {path}\n",
|
|
},
|
|
**kwargs,
|
|
)
|
|
if hash:
|
|
_assert_filehash(path=temp_path, hash=hash, quiet=quiet)
|
|
with filelock.FileLock(lock_path):
|
|
shutil.move(temp_path, path)
|
|
except Exception:
|
|
shutil.rmtree(temp_root)
|
|
raise
|
|
|
|
# postprocess
|
|
if postprocess is not None:
|
|
postprocess(path)
|
|
|
|
return path
|
|
|
|
|
|
def _compute_filehash(path: str, algorithm: str) -> str:
|
|
BLOCKSIZE = 65536
|
|
|
|
if algorithm not in hashlib.algorithms_guaranteed:
|
|
raise ValueError(
|
|
f"Unsupported hash algorithm: {algorithm}. "
|
|
f"Supported algorithms: {hashlib.algorithms_guaranteed}"
|
|
)
|
|
|
|
algorithm_instance = getattr(hashlib, algorithm)()
|
|
with open(path, "rb") as f:
|
|
for block in iter(lambda: f.read(BLOCKSIZE), b""):
|
|
algorithm_instance.update(block)
|
|
return f"{algorithm}:{algorithm_instance.hexdigest()}"
|
|
|
|
|
|
def _assert_filehash(path: str, hash: str, quiet: bool = False) -> None:
|
|
if ":" not in hash:
|
|
raise ValueError(
|
|
f"Invalid hash: {hash}. "
|
|
"Hash must be in the format of {algorithm}:{hash_value}."
|
|
)
|
|
algorithm = hash.split(":")[0]
|
|
|
|
hash_actual = _compute_filehash(path=path, algorithm=algorithm)
|
|
|
|
if hash_actual != hash:
|
|
raise AssertionError(
|
|
f"File hash doesn't match:\nactual: {hash_actual}\nexpected: {hash}"
|
|
)
|