Add a function to handle errors more descriptively

Author: pvk-developer

pyproject.toml

@@ -219,7 +219,6 @@ ignore = [
     # pydocstyle
     "D107",  # Missing docstring in __init__
     "D417",   # Missing argument descriptions in the docstring, this is a bug from pydocstyle: https://github.com/PyCQA/pydocstyle/issues/449
-    "PD901",
     "PD101",
 ]
 

sdv/datasets/demo.py

@@ -6,6 +6,7 @@
 import os
 import warnings
 from collections import defaultdict
+from functools import wraps
 from pathlib import Path
 from zipfile import ZipFile
 
@@ -171,6 +172,56 @@ def is_direct_json_under_prefix(key):
     )
 
 
+def handle_download_failure():
+    """Decorator to handle download exceptions.
+
+    Returns:
+        func:
+            A wrapped function.
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                function_result = func(*args, **kwargs)
+            except DemoResourceNotFoundError as error:
+                bucket = kwargs.get('bucket', args[2])
+                dataset_name = kwargs.get('dataset_name', args[1])
+                error_msg = error.message
+                if 'No objects found under' in error_msg:
+                    raise DemoResourceNotFoundError(
+                        f'Could not download dataset {dataset_name} from bucket {bucket}. '
+                        'Make sure the bucket name is correct. If the bucket is private '
+                        'make sure to provide your credentials.'
+                    ) from error
+
+                if 'METADATA_SPEC_VERSION' in error_msg or 'metadata' in error_msg:
+                    raise DemoResourceNotFoundError(
+                        f'Could not download dataset {dataset_name} from bucket {bucket}. '
+                        'The dataset is missing a valid metadata.'
+                    ) from error
+
+                if 'no csv files were found in data.zip' in error_msg:
+                    raise DemoResourceNotFoundError(
+                        f'Could not download dataset {dataset_name} from bucket {bucket}. '
+                        'The dataset is missing `csv` file/s.'
+                    ) from error
+
+                if "Could not find 'data.zip'" in error_msg:
+                    raise DemoResourceNotFoundError(
+                        f'Could not download dataset {dataset_name} from bucket {bucket}. '
+                        "The dataset is missing 'data.zip' file."
+                    ) from error
+
+            return function_result
+
+        return wrapper
+
+    return decorator
+
+
+@handle_download_failure()
 def _download(modality, dataset_name, bucket, credentials=None):
     """Download dataset resources from a bucket.
 

sdv/errors.py

@@ -104,6 +104,10 @@ class DemoResourceNotFoundError(Exception):
     metadata, license, README, or other auxiliary files in the demo bucket.
     """
 
+    def __init__(self, message):
+        self.message = message
+        super().__init__(self.message)
+
 
 class DemoResourceNotFoundWarning(UserWarning):
     """Warning raised when an optional demo resource is not available.

tests/unit/datasets/test_demo.py

@@ -30,6 +30,7 @@ def _make_zip_with_csv(csv_name: str, df: pd.DataFrame) -> bytes:
     buf = io.BytesIO()
     with zipfile.ZipFile(buf, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
         zf.writestr(csv_name, df.to_csv(index=False))
+
     return buf.getvalue()
 
 
@@ -80,6 +81,7 @@ def side_effect(key, bucket='test_bucket', client=None):
             return zip_bytes
         if key.endswith('metadata.json'):
             return meta_bytes
+
         raise KeyError(key)
 
     mock_get.side_effect = side_effect
@@ -608,16 +610,19 @@ def side_effect(key, bucket, client):
     assert metadata.to_dict()['tables']['word']['primary_key'] == 'id'
 
 
-@patch('sdv.datasets.demo._get_data_from_bucket', return_value=b'{}')
 @patch('sdv.datasets.demo._list_objects')
-def test_download_demo_missing_zip_raises(mock_list, _mock_get):
+def test_download_demo_missing_zip_raises(mock_list):
     # Setup
     mock_list.return_value = [
         {'Key': 'single_table/word/metadata.json'},
     ]
 
     # Run and Assert
-    with pytest.raises(DemoResourceNotFoundError, match="Could not find 'data.zip'"):
+    expected_msg = (
+        'Could not download dataset word from bucket sdv-datasets-public. '
+        "The dataset is missing 'data.zip' file."
+    )
+    with pytest.raises(DemoResourceNotFoundError, match=expected_msg):
         download_demo('single_table', 'word')
 
 
@@ -634,7 +639,11 @@ def test_download_demo_no_v1_metadata_raises(mock_list, mock_get):
     )
 
     # Run and Assert
-    with pytest.raises(DemoResourceNotFoundError, match='METADATA_SPEC_VERSION'):
+    error_msg = (
+        'Could not download dataset word from bucket sdv-datasets-public. '
+        'The dataset is missing a valid metadata.'
+    )
+    with pytest.raises(DemoResourceNotFoundError, match=error_msg):
         download_demo('single_table', 'word')
 
 
@@ -687,8 +696,8 @@ def test_download_demo_writes_metadata_and_discovers_nested_csv(mock_list, mock_
     buf = io.BytesIO()
     with zipfile.ZipFile(buf, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
         zf.writestr('level1/level2/my_table.csv', df.to_csv(index=False))
-    zip_bytes = buf.getvalue()
 
+    zip_bytes = buf.getvalue()
     meta_dict = {
         'METADATA_SPEC_VERSION': 'V1',
         'tables': {