Count the Dl Size progress bar in MiB on every download path

tensorflow_datasets/core/download/downloader.py

@@ -216,8 +216,20 @@ def increase_tqdm(self, url_info: checksums_lib.UrlInfo) -> None:
     """Update the tqdm bars to visually indicate the url_info is downloaded."""
     self._pbar_url.update_total(1)
     self._pbar_url.update(1)
-    self._pbar_dl_size.update_total(url_info.size)
-    self._pbar_dl_size.update(url_info.size)
+    self._update_dl_size_pbar(url_info.size)
+
+  def _update_dl_size_pbar(self, size_in_bytes: int) -> None:
+    """Adds a fully available file of `size_in_bytes` to the `Dl Size` bar.
+
+    The `Dl Size` bar counts MiB (see `tqdm()`), so the size is converted the
+    same way as in `_sync_download`.
+
+    Args:
+      size_in_bytes: Size of the file, in bytes.
+    """
+    size_in_mib = size_in_bytes // units.MiB
+    self._pbar_dl_size.update_total(size_in_mib)
+    self._pbar_dl_size.update(size_in_mib)
 
   def download(
       self, url: str, destination_path: epath.Path, verify: bool = True
@@ -252,8 +264,7 @@ def _sync_file_copy(
     url_info = checksums_lib.compute_url_info(
         out_path, checksum_cls=self._checksumer_cls
     )
-    self._pbar_dl_size.update_total(url_info.size)
-    self._pbar_dl_size.update(url_info.size)
+    self._update_dl_size_pbar(url_info.size)
     self._pbar_url.update(1)
     return DownloadResult(path=out_path, url_info=url_info)
 
@@ -291,10 +302,9 @@ def _sync_download(
       path = destination_path / fname
       size = 0
 
-      # Initialize the download size progress bar
-      size_mb = 0
-      unit_mb = units.MiB
-      total_size = int(response.headers.get('Content-length', 0)) // unit_mb
+      # Initialize the download size progress bar (in MiB)
+      size_mib = 0
+      total_size = int(response.headers.get('Content-length', 0)) // units.MiB
       self._pbar_dl_size.update_total(total_size)
       with path.open('wb') as file_:
         checksum = self._checksumer_cls()
@@ -303,11 +313,11 @@ def _sync_download(
           checksum.update(block)
           file_.write(block)
 
-          # Update the download size progress bar
-          size_mb += len(block)
-          if size_mb > unit_mb:
-            self._pbar_dl_size.update(size_mb // unit_mb)
-            size_mb %= unit_mb
+          # Update the download size progress bar with each completed MiB
+          new_size_mib = size // units.MiB
+          if new_size_mib > size_mib:
+            self._pbar_dl_size.update(new_size_mib - size_mib)
+            size_mib = new_size_mib
     self._pbar_url.update(1)
     return DownloadResult(
         path=path,

tensorflow_datasets/core/download/downloader_test.py

@@ -15,13 +15,17 @@
 
 import hashlib
 import io
+import os
 from typing import Optional
 from unittest import mock
 
 import bs4
 from etils import epath
 import pytest
 from tensorflow_datasets import testing
+from tensorflow_datasets.core import units
+from tensorflow_datasets.core import utils
+from tensorflow_datasets.core.download import checksums
 from tensorflow_datasets.core.download import downloader
 from tensorflow_datasets.core.download import resource as resource_lib
 from tensorflow_datasets.core.download import util
@@ -152,6 +156,44 @@ def test_ftp_error(self):
     with self.assertRaises(downloader.urllib.error.URLError):
       promise.get()
 
+  def _dl_size_pbar_counts(self) -> tuple[int, int]:
+    """Returns the (total, current) MiB counted by the `Dl Size` bar."""
+    pbar = self.downloader._pbar_dl_size
+    total = sum(call.args[0] for call in pbar.update_total.call_args_list)
+    current = sum(call.args[0] for call in pbar.update.call_args_list)
+    return total, current
+
+  def test_dl_size_pbar_cached_file_counts_mib(self):
+    url_info = checksums.UrlInfo(
+        size=utils.Size(3 * units.MiB),
+        checksum=self.resp_checksum,
+        filename='foo.tar.gz',
+    )
+    self.downloader.increase_tqdm(url_info)
+    self.assertEqual(self._dl_size_pbar_counts(), (3, 3))
+
+  def test_dl_size_pbar_file_copy_counts_mib(self):
+    src_path = self.tmp_dir / 'src' / 'foo.tar.gz'
+    src_path.parent.mkdir()
+    src_path.write_bytes(b'x' * (2 * units.MiB))
+    dst_dir = self.tmp_dir / 'dst'
+    dst_dir.mkdir()
+    future = self.downloader.download(os.fspath(src_path), dst_dir).get()
+    self.assertEqual(future.url_info.size, 2 * units.MiB)
+    self.assertEqual(self._dl_size_pbar_counts(), (2, 2))
+
+  def test_dl_size_pbar_http_counts_last_mib(self):
+    content = (b'x' * (64 * 1024 - 1) + b'\n') * 48  # 3 MiB in 64 KiB chunks.
+    mock.patch.object(
+        downloader.requests.Session,
+        'get',
+        lambda *a, **kw: _FakeResponse(
+            self.url, content, headers={'Content-length': len(content)}
+        ),
+    ).start()
+    self.downloader.download(self.url, self.tmp_dir).get()
+    self.assertEqual(self._dl_size_pbar_counts(), (3, 3))
+
 
 # Gramar examples inspired from: https://tools.ietf.org/html/rfc6266#section-5
 _CONTENT_DISPOSITION_FILENAME_PAIRS = [