Merge pull request #195 from ImagingDataCommons/copilot/raise-exception-disk-space

Raise exception when there is not enough disk space

Author: fedorov

idc_index/__init__.py

@@ -9,6 +9,7 @@
 
 __all__ = ["__version__"]
 
-from .index import IDCClient
+from .index import IDCClient, IDCClientInsufficientDiskSpaceError
 
 _ = IDCClient
+_ = IDCClientInsufficientDiskSpaceError

idc_index/cli.py

@@ -11,7 +11,7 @@
 import click
 
 from . import index
-from .index import IDCClient
+from .index import IDCClient, IDCClientInsufficientDiskSpaceError
 
 # Set up logging for the CLI module
 logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.DEBUG)
@@ -190,19 +190,22 @@ def download_from_selection(
     logger_cli.debug(f"use_s5cmd_sync: {use_s5cmd_sync}")
     logger_cli.debug(f"dirTemplate: {dir_template}")
 
-    client.download_from_selection(
-        download_dir,
-        dry_run=dry_run,
-        collection_id=collection_id,
-        patientId=patient_id,
-        studyInstanceUID=study_instance_uid,
-        seriesInstanceUID=series_instance_uid,
-        crdc_series_uuid=crdc_series_uuid,
-        quiet=quiet,
-        show_progress_bar=show_progress_bar,
-        use_s5cmd_sync=use_s5cmd_sync,
-        dirTemplate=dir_template,
-    )
+    try:
+        client.download_from_selection(
+            download_dir,
+            dry_run=dry_run,
+            collection_id=collection_id,
+            patientId=patient_id,
+            studyInstanceUID=study_instance_uid,
+            seriesInstanceUID=series_instance_uid,
+            crdc_series_uuid=crdc_series_uuid,
+            quiet=quiet,
+            show_progress_bar=show_progress_bar,
+            use_s5cmd_sync=use_s5cmd_sync,
+            dirTemplate=dir_template,
+        )
+    except IDCClientInsufficientDiskSpaceError as e:
+        logger_cli.error(e.message)
 
 
 @idc.command()
@@ -286,15 +289,18 @@ def download_from_manifest(
     logger_cli.debug(f"dirTemplate: {dir_template}")
 
     # Call IDCClient's download_from_manifest method with the provided parameters
-    client.download_from_manifest(
-        manifestFile=manifest_file,
-        downloadDir=download_dir,
-        quiet=quiet,
-        validate_manifest=validate_manifest,
-        show_progress_bar=show_progress_bar,
-        use_s5cmd_sync=use_s5cmd_sync,
-        dirTemplate=dir_template,
-    )
+    try:
+        client.download_from_manifest(
+            manifestFile=manifest_file,
+            downloadDir=download_dir,
+            quiet=quiet,
+            validate_manifest=validate_manifest,
+            show_progress_bar=show_progress_bar,
+            use_s5cmd_sync=use_s5cmd_sync,
+            dirTemplate=dir_template,
+        )
+    except IDCClientInsufficientDiskSpaceError as e:
+        logger_cli.error(e.message)
 
 
 @idc.command()
@@ -339,65 +345,68 @@ def download(generic_argument, download_dir, dir_template, log_level):
     else:
         download_dir = Path.cwd()
 
-    if (
-        len(generic_argument) < _get_max_path_length()
-        and Path(generic_argument).is_file()
-    ):
-        # Parse the input parameters and pass them to IDC
-        logger_cli.info("Detected manifest file, downloading from manifest.")
-        client.download_from_manifest(
-            generic_argument, downloadDir=download_dir, dirTemplate=dir_template
-        )
-    # this is not a file manifest
-    else:
-        # Split the input string and filter out any empty values
-        item_ids = [item for item in generic_argument.split(",") if item]
-
-        if not item_ids:
-            logger_cli.error("No valid IDs provided.")
-
-        index_df = client.index
-
-        def check_and_download(column_name, item_ids, download_dir, kwarg_name):
-            matches = index_df[column_name].isin(item_ids)
-            matched_ids = index_df[column_name][matches].unique().tolist()
-            if not matched_ids:
-                return False
-            unmatched_ids = list(set(item_ids) - set(matched_ids))
-            if unmatched_ids:
-                logger_cli.debug(
-                    f"Partial match for {column_name}: matched {matched_ids}, unmatched {unmatched_ids}"
-                )
-            logger_cli.info(f"Identified matching {column_name}: {matched_ids}")
-            client.download_from_selection(
-                **{
-                    kwarg_name: matched_ids,
-                    "downloadDir": download_dir,
-                    "dirTemplate": dir_template,
-                }
+    try:
+        if (
+            len(generic_argument) < _get_max_path_length()
+            and Path(generic_argument).is_file()
+        ):
+            # Parse the input parameters and pass them to IDC
+            logger_cli.info("Detected manifest file, downloading from manifest.")
+            client.download_from_manifest(
+                generic_argument, downloadDir=download_dir, dirTemplate=dir_template
             )
-            return True
+        # this is not a file manifest
+        else:
+            # Split the input string and filter out any empty values
+            item_ids = [item for item in generic_argument.split(",") if item]
+
+            if not item_ids:
+                logger_cli.error("No valid IDs provided.")
+
+            index_df = client.index
+
+            def check_and_download(column_name, item_ids, download_dir, kwarg_name):
+                matches = index_df[column_name].isin(item_ids)
+                matched_ids = index_df[column_name][matches].unique().tolist()
+                if not matched_ids:
+                    return False
+                unmatched_ids = list(set(item_ids) - set(matched_ids))
+                if unmatched_ids:
+                    logger_cli.debug(
+                        f"Partial match for {column_name}: matched {matched_ids}, unmatched {unmatched_ids}"
+                    )
+                logger_cli.info(f"Identified matching {column_name}: {matched_ids}")
+                client.download_from_selection(
+                    **{
+                        kwarg_name: matched_ids,
+                        "downloadDir": download_dir,
+                        "dirTemplate": dir_template,
+                    }
+                )
+                return True
 
-        matches_found = 0
-        matches_found += check_and_download(
-            "collection_id", item_ids, download_dir, "collection_id"
-        )
-        matches_found += check_and_download(
-            "PatientID", item_ids, download_dir, "patientId"
-        )
-        matches_found += check_and_download(
-            "StudyInstanceUID", item_ids, download_dir, "studyInstanceUID"
-        )
-        matches_found += check_and_download(
-            "SeriesInstanceUID", item_ids, download_dir, "seriesInstanceUID"
-        )
-        matches_found += check_and_download(
-            "crdc_series_uuid", item_ids, download_dir, "crdc_series_uuid"
-        )
-        if not matches_found:
-            logger_cli.error(
-                "None of the values passed matched any of the identifiers: collection_id, PatientID, StudyInstanceUID, SeriesInstanceUID, crdc_series_uuid."
+            matches_found = 0
+            matches_found += check_and_download(
+                "collection_id", item_ids, download_dir, "collection_id"
             )
+            matches_found += check_and_download(
+                "PatientID", item_ids, download_dir, "patientId"
+            )
+            matches_found += check_and_download(
+                "StudyInstanceUID", item_ids, download_dir, "studyInstanceUID"
+            )
+            matches_found += check_and_download(
+                "SeriesInstanceUID", item_ids, download_dir, "seriesInstanceUID"
+            )
+            matches_found += check_and_download(
+                "crdc_series_uuid", item_ids, download_dir, "crdc_series_uuid"
+            )
+            if not matches_found:
+                logger_cli.error(
+                    "None of the values passed matched any of the identifiers: collection_id, PatientID, StudyInstanceUID, SeriesInstanceUID, crdc_series_uuid."
+                )
+    except IDCClientInsufficientDiskSpaceError as e:
+        logger_cli.error(e.message)
 
 
 if __name__ == "__main__":

idc_index/index.py

@@ -30,6 +30,30 @@
 logger = logging.getLogger(__name__)
 
 
+class IDCClientInsufficientDiskSpaceError(Exception):
+    """Exception raised when there is insufficient disk space for download."""
+
+    def __init__(
+        self, disk_space_needed: str, disk_space_available: str, message: str = None
+    ):
+        """Initialize the exception.
+
+        Args:
+            disk_space_needed: Human-readable string of disk space needed.
+            disk_space_available: Human-readable string of disk space available.
+            message: Optional custom message. If not provided, a default message is used.
+        """
+        self.disk_space_needed = disk_space_needed
+        self.disk_space_available = disk_space_available
+        if message is None:
+            message = (
+                f"Insufficient disk space. "
+                f"Need {disk_space_needed}, but only {disk_space_available} available."
+            )
+        self.message = message
+        super().__init__(self.message)
+
+
 class IDCClient:
     # Default download hierarchy template
     DOWNLOAD_HIERARCHY_DEFAULT = (
@@ -314,14 +338,27 @@ def _check_create_directory(download_dir):
 
         return str(download_dir.resolve())
 
-    def _check_disk_size_and_warn(self, download_dir, disk_size_needed):
+    def _validate_disk_space(self, download_dir, disk_size_needed):
+        """Check if there is sufficient disk space for the download.
+
+        Args:
+            download_dir: The directory where files will be downloaded.
+            disk_size_needed: The size needed in MB.
+
+        Raises:
+            IDCClientInsufficientDiskSpaceError: If there is not enough disk space.
+        """
         disk_free_space_MB = psutil.disk_usage(download_dir).free / (1000 * 1000)
-        logger.info("Disk size needed: " + self._format_size(disk_size_needed))
-        logger.info("Disk size available: " + self._format_size(disk_free_space_MB))
+        disk_size_needed_str = self._format_size(disk_size_needed)
+        disk_free_space_str = self._format_size(disk_free_space_MB)
+        logger.info("Disk size needed: " + disk_size_needed_str)
+        logger.info("Disk size available: " + disk_free_space_str)
         if disk_free_space_MB < disk_size_needed:
             logger.error("Not enough free space on disk to download the files.")
-            return False
-        return True
+            raise IDCClientInsufficientDiskSpaceError(
+                disk_space_needed=disk_size_needed_str,
+                disk_space_available=disk_free_space_str,
+            )
 
     def fetch_index(self, index_name) -> None:
         """Downloads requested index and adds this index joined with the main index as respective class attribute.
@@ -1591,6 +1628,7 @@ def download_from_manifest(
 
         Raises:
             ValueError: If the download directory does not exist.
+            IDCClientInsufficientDiskSpaceError: If there is not enough disk space.
         """
         downloadDir = self._check_create_directory(downloadDir)
 
@@ -1610,8 +1648,7 @@ def download_from_manifest(
         )
 
         total_size_rounded = round(total_size, 2)
-        if not self._check_disk_size_and_warn(downloadDir, total_size):
-            return
+        self._validate_disk_space(downloadDir, total_size)
 
         self._s5cmd_run(
             endpoint_to_use=endpoint_to_use,
@@ -1841,8 +1878,7 @@ def download_from_selection(
             total_size_bytes = round(result_df["instance_size"].sum(), 2)
             total_size = total_size_bytes / (10**6)
 
-        if not self._check_disk_size_and_warn(downloadDir, total_size):
-            return
+        self._validate_disk_space(downloadDir, total_size)
 
         if dry_run:
             logger.info(

tests/idcindex.py

@@ -6,12 +6,13 @@
 import unittest
 from itertools import product
 from pathlib import Path
+from unittest.mock import patch
 
 import pandas as pd
 import pytest
 import requests
 from click.testing import CliRunner
-from idc_index import IDCClient, cli
+from idc_index import IDCClient, IDCClientInsufficientDiskSpaceError, cli
 
 # Run tests using the following command from the root of the repository:
 # python -m unittest -vv tests/idcindex.py
@@ -604,5 +605,71 @@ def test_instance_file_URLs(self):
         assert files_aws == files_gcp == file_url
 
 
+class TestInsufficientDiskSpaceException(unittest.TestCase):
+    def setUp(self):
+        self.client = IDCClient()
+
+    @staticmethod
+    def _create_mock_disk_usage(free_bytes=1000):
+        """Create a mock disk usage object with the specified free space."""
+        return type("DiskUsage", (), {"free": free_bytes})()
+
+    def test_exception_attributes(self):
+        """Test that the exception has the correct attributes."""
+        exc = IDCClientInsufficientDiskSpaceError(
+            disk_space_needed="10 GB",
+            disk_space_available="5 GB",
+        )
+        assert exc.disk_space_needed == "10 GB"
+        assert exc.disk_space_available == "5 GB"
+        assert "10 GB" in str(exc)
+        assert "5 GB" in str(exc)
+        assert "Insufficient disk space" in str(exc)
+
+    def test_exception_custom_message(self):
+        """Test that a custom message can be provided."""
+        custom_msg = "Custom error message"
+        exc = IDCClientInsufficientDiskSpaceError(
+            disk_space_needed="10 GB",
+            disk_space_available="5 GB",
+            message=custom_msg,
+        )
+        assert exc.message == custom_msg
+        assert str(exc) == custom_msg
+
+    def test_exception_raised_on_insufficient_space(self):
+        """Test that exception is raised when disk space is insufficient."""
+        # Mock the disk check to simulate insufficient space (1000 bytes = ~0.001 MB)
+        mock_usage = self._create_mock_disk_usage(free_bytes=1000)
+        with tempfile.TemporaryDirectory() as temp_dir:
+            with patch("psutil.disk_usage", return_value=mock_usage):
+                with pytest.raises(IDCClientInsufficientDiskSpaceError) as exc_info:
+                    self.client.download_from_selection(
+                        downloadDir=temp_dir,
+                        seriesInstanceUID="1.3.6.1.4.1.14519.5.2.1.7695.1700.153974929648969296590126728101",
+                    )
+                assert "Insufficient disk space" in str(exc_info.value)
+
+    def test_cli_handles_insufficient_space_gracefully(self):
+        """Test that CLI handles the exception without crashing."""
+        # Mock the disk check to simulate insufficient space (1000 bytes = ~0.001 MB)
+        mock_usage = self._create_mock_disk_usage(free_bytes=1000)
+        runner = CliRunner()
+        with patch("psutil.disk_usage", return_value=mock_usage):
+            result = runner.invoke(
+                cli.download_from_selection,
+                [
+                    "--download-dir",
+                    "/tmp",
+                    "--study-instance-uid",
+                    "1.3.6.1.4.1.14519.5.2.1.7695.1700.114861588187429958687900856462",
+                ],
+            )
+            # The CLI should not crash (exit code 0) - the error is logged
+            assert result.exit_code == 0
+            # The exception was raised, handled, and logged (not re-raised)
+            assert result.exception is None
+
+
 if __name__ == "__main__":
     unittest.main()