Added support for puling image from an URL to upload_image.py

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

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 <name_of_example>
 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 <name_of_stack> 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__":