Show progress while upload computes checksums

mapillary_tools/uploader.py

@@ -32,6 +32,7 @@
     from typing_extensions import Required
 
 import requests
+from tqdm import tqdm
 
 from . import (
     api_v4,
@@ -248,7 +249,10 @@ def upload_videos(
         for idx, video_metadata in enumerate(sorted_video_metadatas):
             LOG.debug(f"Checksum for video {video_metadata.filename}...")
             try:
-                video_metadata.update_md5sum()
+                _update_md5sums(
+                    [video_metadata],
+                    desc=f"Checksumming {video_metadata.filetype.value.upper()} {video_metadata.filename.name} ({idx + 1}/{len(sorted_video_metadatas)})",
+                )
             except Exception as ex:
                 yield video_metadata, UploadResult(error=ex)
                 continue
@@ -603,6 +607,10 @@ def upload_images(
 
         for sequence_idx, (sequence_uuid, sequence) in enumerate(sequences.items()):
             LOG.debug(f"Checksum for image sequence {sequence_uuid}...")
+            _update_md5sums(
+                sequence,
+                desc=f"Checksumming IMAGE ({sequence_idx + 1}/{len(sequences)})",
+            )
             sequence_md5sum = types.update_sequence_md5sum(sequence)
 
             sequence_progress: SequenceProgress = {
@@ -1117,6 +1125,32 @@ def _gen_session_key(self, fp: T.IO[bytes], progress: dict[str, T.Any]) -> str:
         return session_key
 
 
+def _update_md5sums(
+    metadatas: T.Sequence[types.ImageMetadata] | T.Sequence[types.VideoMetadata],
+    desc: str,
+) -> None:
+    """
+    Compute the checksums that the metadatas miss, showing the progress, as
+    reading the files takes long for large sequences or slow drives, before
+    the upload history is checked.
+    """
+    missing = [metadata for metadata in metadatas if metadata.md5sum is None]
+    if not missing:
+        return
+
+    with tqdm(
+        total=sum(utils.get_file_size(metadata.filename) for metadata in missing),
+        desc=desc,
+        unit="B",
+        unit_scale=True,
+        unit_divisor=1024,
+        disable=LOG.isEnabledFor(logging.DEBUG),
+    ) as pbar:
+        for metadata in missing:
+            with metadata.filename.open("rb") as fp:
+                metadata.md5sum = utils.md5sum_fp(fp, progress=pbar.update).hexdigest()
+
+
 def _validate_metadatas(metadatas: T.Sequence[types.ImageMetadata]):
     for metadata in metadatas:
         validate_image_desc(DescriptionJSONSerializer.as_desc(metadata))

mapillary_tools/utils.py

@@ -15,14 +15,24 @@
 
 # Use "hashlib._Hash" instead of hashlib._Hash because:
 # AttributeError: module 'hashlib' has no attribute '_Hash'
-def md5sum_fp(fp: T.IO[bytes], md5: "hashlib._Hash | None" = None) -> "hashlib._Hash":
+def md5sum_fp(
+    fp: T.IO[bytes],
+    md5: "hashlib._Hash | None" = None,
+    progress: T.Callable[[int], T.Any] | None = None,
+) -> "hashlib._Hash":
+    """
+    Update the hash with the content of the file, calling progress with the
+    number of bytes read after each chunk.
+    """
     if md5 is None:
         md5 = hashlib.md5()
     while True:
         buf = fp.read(1024 * 1024 * 32)
         if not buf:
             break
         md5.update(buf)
+        if progress is not None:
+            progress(len(buf))
     return md5
 
 

tests/unit/test_uploader.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import dataclasses
+import hashlib
 import typing as T
 from pathlib import Path
 from unittest.mock import patch
@@ -824,3 +825,30 @@ def on_upload_finished(payload):
         )
         assert "entity_size" in start_payload, "upload_start should contain entity_size"
         assert start_payload["sequence_uuid"] == "event_test_sequence"
+
+
+def test_update_md5sums(tmpdir: py.path.local):
+    image_path = Path(tmpdir.join("image.jpg"))
+    image_path.write_bytes(b"image bytes")
+    other_path = Path(tmpdir.join("other.jpg"))
+    other_path.write_bytes(b"other bytes")
+    metadatas = [
+        types.ImageMetadata(
+            filename=image_path, lat=0, lon=0, alt=None, angle=None, time=0
+        ),
+        types.ImageMetadata(
+            filename=other_path,
+            lat=0,
+            lon=0,
+            alt=None,
+            angle=None,
+            time=1,
+            md5sum="0" * 32,
+        ),
+    ]
+
+    uploader._update_md5sums(metadatas, desc="Checksumming IMAGE (1/1)")
+
+    assert metadatas[0].md5sum == hashlib.md5(b"image bytes").hexdigest()
+    # Checksums given are kept
+    assert metadatas[1].md5sum == "0" * 32

tests/unit/test_utils.py

@@ -3,6 +3,8 @@
 # This source code is licensed under the BSD license found in the
 # LICENSE file in the root directory of this source tree.
 
+import hashlib
+import io
 import logging
 import multiprocessing
 import os
@@ -344,3 +346,13 @@ def test_empty_after_sanitize_returns_none(self):
     def test_preserves_case(self):
         """Test that case is preserved"""
         assert utils.sanitize_serial("AbC123xYz") == "AbC123xYz"
+
+
+def test_md5sum_fp_progress():
+    data = b"x" * 1000
+    progress: list[int] = []
+
+    md5 = utils.md5sum_fp(io.BytesIO(data), progress=progress.append)
+
+    assert md5.hexdigest() == hashlib.md5(data).hexdigest()
+    assert sum(progress) == len(data)