Address failing _check_disk_space() when path doesn't exist yet (#1692)

* Add realpath to disk_usage to allow path-like obj * Add logic to address non-existent paths * Add test cased for non-existent paths * Update tests/test_file_download.py --------- Co-authored-by: Lucain <[email protected]>
huggingface · Sep 26, 2023 · 89cc691 · 89cc691
1 parent fe711bd
commit 89cc691
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 11 deletions.
diff --git a/src/huggingface_hub/file_download.py b/src/huggingface_hub/file_download.py
@@ -972,17 +972,19 @@ def _check_disk_space(expected_size: int, target_dir: Union[str, Path]) -> None:
             The directory where the file will be stored after downloading.
     """
 
-    target_dir = str(target_dir)
-    target_dir_free = shutil.disk_usage(target_dir).free
-
-    has_enough_space = target_dir_free >= expected_size
-
-    if not has_enough_space:
-        warnings.warn(
-            "Not enough free disk space to download the file. "
-            f"The expected file size is: {expected_size / 1e6:.2f} MB. "
-            f"The target location {target_dir} only has {target_dir_free / 1e6:.2f} MB free disk space."
-        )
+    target_dir = Path(target_dir)  # format as `Path`
+    for path in [target_dir] + list(target_dir.parents):  # first check target_dir, then each parents one by one
+        try:
+            target_dir_free = shutil.disk_usage(path).free
+            if target_dir_free < expected_size:
+                warnings.warn(
+                    "Not enough free disk space to download the file. "
+                    f"The expected file size is: {expected_size / 1e6:.2f} MB. "
+                    f"The target location {target_dir} only has {target_dir_free / 1e6:.2f} MB free disk space."
+                )
+            return
+        except OSError:  # raise on anything: file does not exist or space disk cannot be checked
+            pass
 
 
 @validate_hf_hub_args

diff --git a/tests/test_file_download.py b/tests/test_file_download.py
@@ -112,6 +112,21 @@ def test_disk_usage_warning(self, disk_usage_mock: Mock) -> None:
             _check_disk_space(expected_size=self.expected_size, target_dir=disk_usage_mock)
             assert len(w) == 0
 
+    def test_disk_usage_warning_with_non_existent_path(self) -> None:
+        # Test for not existent (absolute) path
+        with warnings.catch_warnings(record=True) as w:
+            # Cause all warnings to always be triggered.
+            warnings.simplefilter("always")
+            _check_disk_space(expected_size=self.expected_size, target_dir="path/to/not_existent_path")
+            assert len(w) == 0
+
+        # Test for not existent (relative) path
+        with warnings.catch_warnings(record=True) as w:
+            # Cause all warnings to always be triggered.
+            warnings.simplefilter("always")
+            _check_disk_space(expected_size=self.expected_size, target_dir="/path/to/not_existent_path")
+            assert len(w) == 0
+
 
 @with_production_testing
 class CachedDownloadTests(unittest.TestCase):