diff --git a/scripts/requirements-dev.txt b/scripts/requirements-dev.txt index f1a3c27ad3c3..91c041c08ddc 100644 --- a/scripts/requirements-dev.txt +++ b/scripts/requirements-dev.txt @@ -8,3 +8,5 @@ cryptography==38.0.4 # for scripts/upload_image.py google-cloud-storage==2.9.0 # for scripts/upload_image.py PyGithub==1.58.2 # for scripts/generate_pr_summary.py and scripts/pr_link_docs_preview.py Pillow # for scripts/upload_image.py +tqdm +requests diff --git a/scripts/upload_image.py b/scripts/upload_image.py index 27b72f90c1be..23afe61484e9 100755 --- a/scripts/upload_image.py +++ b/scripts/upload_image.py @@ -45,12 +45,16 @@ import subprocess import sys import tempfile +import urllib.parse +import urllib.request from io import BytesIO from pathlib import Path import PIL import PIL.Image import PIL.ImageGrab +import requests +import tqdm from google.cloud import storage from PIL.Image import Image, Resampling @@ -312,6 +316,45 @@ def data_hash(data: bytes) -> str: return hashlib.sha1(data).hexdigest() +def download_file(url: str, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + logging.info("Downloading %s to %s", url, path) + response = requests.get(url, stream=True) + with tqdm.tqdm.wrapattr( + open(path, "wb"), + "write", + miniters=1, + total=int(response.headers.get("content-length", 0)), + desc=f"Downloading {path.name}", + ) as f: + for chunk in response.iter_content(chunk_size=4096): + f.write(chunk) + + +def run(args) -> None: + """Run the script based on the provided args.""" + try: + if shutil.which("pngcrush") is None and not args.skip_pngcrush: + raise RuntimeError("pngcrush is not installed, consider using --skip-pngcrush") + + uploader = Uploader(not args.skip_pngcrush) + + if args.single: + object_name = uploader.upload_file(args.path) + print(f"\nhttps://static.rerun.io/{object_name}") + else: + if args.path is None: + if args.name is None: + raise RuntimeError("Name is required when uploading from clipboard") + else: + html_str = uploader.upload_stack_from_clipboard(args.name) + else: + html_str = uploader.upload_stack_from_file(args.path, args.name) + print("\n" + html_str) + except RuntimeError as e: + print(f"Error: {e.args[0]}", file=sys.stderr) + + DESCRIPTION = """Upload an image to static.rerun.io. Example screenshots @@ -324,13 +367,20 @@ def data_hash(data: bytes) -> str: 4. Take a screenshot using the command palette. 5. Run: just upload --name 6. Copy the output HTML tag and paste it into the README.md file. + +Other uses +---------- + +Download an image, optimize it and create a multi-resolution stack: + + just upload --name https://example.com/path/to/image.png """ def main() -> None: parser = argparse.ArgumentParser(description=DESCRIPTION, formatter_class=argparse.RawTextHelpFormatter) parser.add_argument( - "path", type=Path, nargs="?", help="File path to the image. If not provided, use the clipboard's content." + "path", type=str, nargs="?", help="Image file URL or path. If not provided, use the clipboard's content." ) parser.add_argument( "--single", action="store_true", help="Upload a single image instead of creating a multi-resolution stack." @@ -345,26 +395,20 @@ def main() -> None: else: logging.basicConfig(level=logging.INFO) - try: - if shutil.which("pngcrush") is None and not args.skip_pngcrush: - raise RuntimeError("pngcrush is not installed, consider using --skip-pngcrush") - - uploader = Uploader(not args.skip_pngcrush) - - if args.single: - object_name = uploader.upload_file(args.path) - print(f"\nhttps://static.rerun.io/{object_name}") - else: - if args.path is None: - if args.name is None: - raise RuntimeError("Name is required when uploading from clipboard") - else: - html_str = uploader.upload_stack_from_clipboard(args.name) + # The entire block is wrapped around tmp_dir such that it exists for the entire run. + with tempfile.TemporaryDirectory() as tmp_dir: + # check if path as a URL and download it. + if args.path is not None: + res = urllib.parse.urlparse(args.path) + if res.scheme and res.netloc: + file_name = os.path.basename(res.path) + local_path = Path(tmp_dir) / file_name + download_file(args.path, local_path) + args.path = Path(local_path) else: - html_str = uploader.upload_stack_from_file(args.path, args.name) - print("\n" + html_str) - except RuntimeError as e: - print(f"Error: {e.args[0]}", file=sys.stderr) + args.path = Path(args.path) + + run(args) if __name__ == "__main__":