diff --git a/tensorflow_datasets/core/download/downloader.py b/tensorflow_datasets/core/download/downloader.py index ee9f1d2ace6..843d636f1e9 100644 --- a/tensorflow_datasets/core/download/downloader.py +++ b/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, diff --git a/tensorflow_datasets/core/download/downloader_test.py b/tensorflow_datasets/core/download/downloader_test.py index 91b534821d5..dc280345201 100644 --- a/tensorflow_datasets/core/download/downloader_test.py +++ b/tensorflow_datasets/core/download/downloader_test.py @@ -15,6 +15,7 @@ import hashlib import io +import os from typing import Optional from unittest import mock @@ -22,6 +23,9 @@ 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 = [