Feat/granules volumetry

altimetry_downloader_aviso/catalog_client/_granules_utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import concurrent.futures as cf
 import logging
 import typing as tp
 import warnings
@@ -8,6 +9,7 @@
 from typing import TYPE_CHECKING
 from urllib.parse import urljoin
 
+import requests
 import yaml
 from fcollections.core import (
     FileNode,
@@ -219,3 +221,44 @@ def _parse_tds_layout(product: AvisoProduct) -> ProductLayoutConfig:
             catalog_path=product_layout["catalog_path"],
             default_filters=product_layout["filters"],
         )
+
+
+def _get_size_from_url(url: str, timeout: float = 5.0) -> tp.Optional[int]:
+    """Return the size in bytes of a remote file via an HTTP HEAD request, or
+    None if unavailable."""
+    try:
+        resp = requests.head(url, allow_redirects=True, timeout=timeout)
+        resp.raise_for_status()
+        content_length = resp.headers.get("Content-Length")
+        return int(content_length) if content_length is not None else None
+    except requests.RequestException as e:
+        logger.warning("Cannot retrieve size for %s: %s", url, e)
+        return None
+
+
+def estimate_total_size(
+    urls: tp.Sequence[str], max_workers: int = 8
+) -> tp.Tuple[int, int]:
+    """Estimate total download size (bytes) for a list of granule URLs.
+
+    Returns (total_bytes, unknown_count), the latter being the number of
+    granules whose size could not be determined via HEAD.
+    """
+    total = 0
+    unknown = 0
+    with cf.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        for size in executor.map(_get_size_from_url, urls):
+            if size is None:
+                unknown += 1
+            else:
+                total += size
+    return total, unknown
+
+
+def format_size(num_bytes: float) -> str:
+    """Format a byte count as a human-readable string (e.g. '1.3 GB')."""
+    for unit in ("B", "KB", "MB", "GB", "TB"):
+        if abs(num_bytes) < 1024:
+            return f"{num_bytes:.1f} {unit}"
+        num_bytes /= 1024
+    return f"{num_bytes:.1f} PB"

altimetry_downloader_aviso/cli.py

@@ -248,6 +248,12 @@ def get(
         "-V",
         help="Product's version. By default, last version is selected",
     ),
+    yes: bool = typer.Option(
+        False,
+        "--yes",
+        "-y",
+        help="Skip the download size confirmation prompt",
+    ),
     quiet: bool = typer.Option(
         False,
         "--quiet",
@@ -278,6 +284,7 @@ def get(
             time=(start, end),
             version=version,
             overwrite=overwrite,
+            assume_yes=yes,
         )
 
         console.print(f"[green]Local files ({len(downloaded_files)}) :[/]")

altimetry_downloader_aviso/core.py

@@ -8,6 +8,7 @@
 import yaml
 
 from .auth import ensure_credentials
+from .catalog_client._granules_utils import estimate_total_size, format_size
 from .catalog_client.client import (
     fetch_catalog,
     get_details,
@@ -69,6 +70,31 @@ def details(product_short_name: str) -> AvisoProduct:
     return get_details(product_short_name)
 
 
+def confirm_download(urls: tp.Sequence[str], assume_yes: bool = False) -> bool:
+    """Print the estimated total volumetry and ask for user confirmation.
+
+    Returns True if the download should proceed (user confirmed, or
+    assume_yes was passed).
+    """
+    if not urls:
+        return True
+
+    total, unknown = estimate_total_size(urls)
+    msg = (
+        f"About to download {len(urls)} file(s),"
+        f"estimated total size: {format_size(total)}"
+    )
+    if unknown:
+        msg += f" ({unknown} size(s) could not be determined)"
+    print(msg)
+
+    if assume_yes:
+        return True
+
+    answer = input("Proceed with download? [y/N] ").strip().lower()
+    return answer in ("y", "yes")
+
+
 @authenticate
 def get(
     product_short_name: str,
@@ -78,6 +104,7 @@ def get(
     time: tuple[np.datetime64, np.datetime64] | None = None,
     version: str | None = None,
     overwrite: bool = False,
+    assume_yes: bool = False,
 ) -> list[str]:
     """Downloads a product from Aviso's Thredds Data Server.
 
@@ -97,6 +124,8 @@ def get(
         the version for files/folders selection
     overwrite: bool
         whether to overwrite files if they already exist
+    assume_yes: bool
+        whether to skip the download confirmation prompt (default: False)
 
     Returns
     -------
@@ -116,6 +145,9 @@ def get(
     granule_paths, _, non_target_local_files = _search_granules_with_overwrite(
         product_short_name, Protocol.HTTP, output_dir, overwrite, **filters
     )
+    if not confirm_download(granule_paths, assume_yes=assume_yes):
+        logger.info("Download cancelled by user.")
+        return non_target_local_files
 
     logger.debug("Downloading granules: %s...", list(granule_paths))
 

tests/conftest.py

@@ -270,6 +270,11 @@ def patch_all(mocker):
         mocker.patch(x)
 
 
+@pytest.fixture(autouse=True)
+def bypass_download_confirmation(mocker):
+    mocker.patch("altimetry_downloader_aviso.core.confirm_download", return_value=True)
+
+
 # PATCH TDS CATALOG CONTENT
 
 

tests/test_cli.py

@@ -199,6 +199,7 @@ def test_get_simple_filters(mocker, tmp_path):
         pass_number=None,
         time=(None, None),
         overwrite=False,
+        assume_yes=False,
     )
 
 

tests/test_core.py

@@ -6,7 +6,13 @@
 
 from altimetry_downloader_aviso.auth import AuthenticationError
 from altimetry_downloader_aviso.catalog_client.client import InvalidProductError
-from altimetry_downloader_aviso.core import details, get, subset, summary
+from altimetry_downloader_aviso.core import (
+    confirm_download,
+    details,
+    get,
+    subset,
+    summary,
+)
 
 
 def test_summary():
@@ -79,6 +85,21 @@ def test_get_subset(tmp_path, short_name, filters, files, command):
     assert local_files == [str(tmp_path / f) for f in files]
 
 
+def test_get_download_cancelled(mocker, tmp_path):
+    mock_confirm = mocker.patch(
+        "altimetry_downloader_aviso.core.confirm_download", return_value=False
+    )
+    with patch("altimetry_downloader_aviso.subset.subset_one_file", return_value=True):
+        local_files = get(
+            product_short_name="sample_product_a",
+            output_dir=tmp_path,
+            cycle_number=2,
+        )
+
+    assert local_files == []
+    mock_confirm.assert_called_once()
+
+
 def test_subset_parameters_passed(tmp_path):
     with patch(
         "altimetry_downloader_aviso.subset.subset_one_file", return_value=True
@@ -167,3 +188,43 @@ def test_get_subset_auth_error(mocker, tmp_path, command):
 @pytest.mark.parametrize("command", [get, subset])
 def test_get_subset_bad_filters(tmp_path, short_name, filters, command):
     assert command(short_name, tmp_path, **filters) == []
+
+
+def test_confirm_download_empty_urls():
+    assert confirm_download([]) is True
+
+
+def test_confirm_download_assume_yes(mocker, capsys):
+    mocker.patch(
+        "altimetry_downloader_aviso.core.estimate_total_size",
+        return_value=(1024, 0),
+    )
+    result = confirm_download(["https://tds.mock/a.nc"], assume_yes=True)
+
+    assert result is True
+    assert "1.0 KB" in capsys.readouterr().out
+
+
+def test_confirm_download_prompt_yes(mocker, capsys):
+    mocker.patch(
+        "altimetry_downloader_aviso.core.estimate_total_size",
+        return_value=(2048, 1),
+    )
+    mocker.patch("builtins.input", return_value="y")
+
+    result = confirm_download(["https://tds.mock/a.nc", "https://tds.mock/b.nc"])
+
+    out = capsys.readouterr().out
+    assert result is True
+    assert "2 file(s)" in out
+    assert "1 size(s) could not be determined" in out
+
+
+def test_confirm_download_prompt_no(mocker):
+    mocker.patch(
+        "altimetry_downloader_aviso.core.estimate_total_size",
+        return_value=(1024, 0),
+    )
+    mocker.patch("builtins.input", return_value="n")
+
+    assert confirm_download(["https://tds.mock/a.nc"]) is False

tests/test_granules_utils.py

@@ -0,0 +1,128 @@
+import requests
+
+from altimetry_downloader_aviso.catalog_client._granules_utils import (
+    _get_size_from_url,
+    estimate_total_size,
+    format_size,
+)
+
+# --- _get_size_from_url ---
+
+
+def test_get_size_from_url_with_content_length(mocker):
+    mock_response = mocker.Mock()
+    mock_response.headers = {"Content-Length": "12345"}
+    mocker.patch(
+        "altimetry_downloader_aviso.catalog_client._granules_utils.requests.head",
+        return_value=mock_response,
+    )
+
+    assert _get_size_from_url("https://tds.mock/a.nc") == 12345
+
+
+def test_get_size_from_url_without_content_length(mocker):
+    mock_response = mocker.Mock()
+    mock_response.headers = {}
+    mocker.patch(
+        "altimetry_downloader_aviso.catalog_client._granules_utils.requests.head",
+        return_value=mock_response,
+    )
+
+    assert _get_size_from_url("https://tds.mock/a.nc") is None
+
+
+def test_get_size_from_url_http_error(mocker):
+    mock_response = mocker.Mock()
+    mock_response.raise_for_status.side_effect = requests.HTTPError("404")
+    mocker.patch(
+        "altimetry_downloader_aviso.catalog_client._granules_utils.requests.head",
+        return_value=mock_response,
+    )
+
+    assert _get_size_from_url("https://tds.mock/a.nc") is None
+
+
+def test_get_size_from_url_connection_error(mocker):
+    mocker.patch(
+        "altimetry_downloader_aviso.catalog_client._granules_utils.requests.head",
+        side_effect=requests.ConnectionError("unreachable"),
+    )
+
+    assert _get_size_from_url("https://tds.mock/a.nc") is None
+
+
+def test_get_size_from_url_passes_timeout(mocker):
+    mock_head = mocker.patch(
+        "altimetry_downloader_aviso.catalog_client._granules_utils.requests.head",
+        return_value=mocker.Mock(headers={"Content-Length": "10"}),
+    )
+
+    _get_size_from_url("https://tds.mock/a.nc", timeout=2.5)
+
+    assert mock_head.call_args.kwargs["timeout"] == 2.5
+
+
+# --- estimate_total_size ---
+
+
+def test_estimate_total_size_empty():
+    assert estimate_total_size([]) == (0, 0)
+
+
+def test_estimate_total_size_all_known(mocker):
+    mocker.patch(
+        "altimetry_downloader_aviso.catalog_client._granules_utils."
+        "_get_size_from_url",
+        side_effect=[100, 200, 300],
+    )
+
+    total, unknown = estimate_total_size(
+        ["https://tds.mock/a.nc", "https://tds.mock/b.nc", "https://tds.mock/c.nc"]
+    )
+
+    assert total == 600
+    assert unknown == 0
+
+
+def test_estimate_total_size_mixed_unknown(mocker):
+    mocker.patch(
+        "altimetry_downloader_aviso.catalog_client._granules_utils."
+        "_get_size_from_url",
+        side_effect=[100, None, 300, None],
+    )
+
+    total, unknown = estimate_total_size(
+        [
+            "https://tds.mock/a.nc",
+            "https://tds.mock/b.nc",
+            "https://tds.mock/c.nc",
+            "https://tds.mock/d.nc",
+        ]
+    )
+
+    assert total == 400
+    assert unknown == 2
+
+
+def test_estimate_total_size_all_unknown(mocker):
+    mocker.patch(
+        "altimetry_downloader_aviso.catalog_client._granules_utils."
+        "_get_size_from_url",
+        return_value=None,
+    )
+
+    total, unknown = estimate_total_size(["https://tds.mock/a.nc"])
+
+    assert total == 0
+    assert unknown == 1
+
+
+# --- format_size ---
+
+
+def test_format_size_bytes():
+    assert format_size(500) == "500.0 B"
+
+
+def test_format_size_petabytes():
+    assert format_size(1024**5 * 1.1) == "1.1 PB"