Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 13 additions & 11 deletions src/huggingface_hub/file_download.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 15 additions & 0 deletions tests/test_file_download.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down