diff --git a/download_imageset.py b/download_imageset.py index ebe53cd..6c804be 100755 --- a/download_imageset.py +++ b/download_imageset.py @@ -391,6 +391,9 @@ def filehost_cmd(urls: tuple[str, ...], name: str | None, cwd: Path, token: str) _extract_and_process(output_dir) +def is_filehost(url: str) -> bool: + kws = ["rg.to", "rapidgator.net", "trbt.cc"] + return any(kw in url for kw in kws) @cli.command("download") @click.argument("urls", nargs=-1, required=True) @@ -401,7 +404,7 @@ def filehost_cmd(urls: tuple[str, ...], name: str | None, cwd: Path, token: str) def download_cmd( ctx: click.Context, urls: tuple[str, ...], name: str | None, cwd: Path, token: str ) -> None: - filehost_urls = [u for u in urls if "rg.to" in u or "rapidgator.net" in u] + filehost_urls = [u for u in urls if is_filehost(u)] gallery_urls = [u for u in urls if u not in set(filehost_urls)] if filehost_urls: ctx.invoke(filehost_cmd, urls=tuple(filehost_urls), name=name, cwd=cwd, token=token)