Pass through additional arguments from HfApi download utils (#3531)

Co-authored-by: célina <[email protected]>
Co-authored-by: Lucain <[email protected]>

Author: 3 people

src/huggingface_hub/hf_api.py

@@ -5216,6 +5216,7 @@ def hf_hub_download(
         etag_timeout: float = constants.DEFAULT_ETAG_TIMEOUT,
         token: Union[bool, str, None] = None,
         local_files_only: bool = False,
+        tqdm_class: Optional[type[base_tqdm]] = None,
         dry_run: Literal[False] = False,
     ) -> str: ...
 
@@ -5234,6 +5235,7 @@ def hf_hub_download(
         etag_timeout: float = constants.DEFAULT_ETAG_TIMEOUT,
         token: Union[bool, str, None] = None,
         local_files_only: bool = False,
+        tqdm_class: Optional[type[base_tqdm]] = None,
         dry_run: Literal[True],
     ) -> DryRunFileInfo: ...
 
@@ -5252,6 +5254,7 @@ def hf_hub_download(
         etag_timeout: float = constants.DEFAULT_ETAG_TIMEOUT,
         token: Union[bool, str, None] = None,
         local_files_only: bool = False,
+        tqdm_class: Optional[type[base_tqdm]] = None,
         dry_run: bool = False,
     ) -> Union[str, DryRunFileInfo]:
         """Download a given file if it's not already present in the local cache.
@@ -5320,6 +5323,11 @@ def hf_hub_download(
             local_files_only (`bool`, *optional*, defaults to `False`):
                 If `True`, avoid downloading the file and return the path to the
                 local cached file if it exists.
+            tqdm_class (`tqdm`, *optional*):
+                If provided, overwrites the default behavior for the progress bar. Passed
+                argument must inherit from `tqdm.auto.tqdm` or at least mimic its behavior.
+                Defaults to the custom HF progress bar that can be disabled by setting
+                `HF_HUB_DISABLE_PROGRESS_BARS` environment variable.
             dry_run (`bool`, *optional*, defaults to `False`):
                 If `True`, perform a dry run without actually downloading the file. Returns a
                 [`DryRunFileInfo`] object containing information about what would be downloaded.
@@ -5369,6 +5377,8 @@ def hf_hub_download(
             token=token,
             headers=self.headers,
             local_files_only=local_files_only,
+            tqdm_class=tqdm_class,
+            dry_run=dry_run,
         )
 
     @validate_hf_hub_args
@@ -5388,7 +5398,8 @@ def snapshot_download(
         ignore_patterns: Optional[Union[list[str], str]] = None,
         max_workers: int = 8,
         tqdm_class: Optional[type[base_tqdm]] = None,
-    ) -> str:
+        dry_run: bool = False,
+    ) -> Union[str, list[DryRunFileInfo]]:
         """Download repo files.
 
         Download a whole snapshot of a repo's files at the specified revision. This is useful when you want all files from
@@ -5443,9 +5454,14 @@ def snapshot_download(
                 Note that the `tqdm_class` is not passed to each individual download.
                 Defaults to the custom HF progress bar that can be disabled by setting
                 `HF_HUB_DISABLE_PROGRESS_BARS` environment variable.
+            dry_run (`bool`, *optional*, defaults to `False`):
+                If `True`, perform a dry run without actually downloading the files. Returns a list of
+                [`DryRunFileInfo`] objects containing information about what would be downloaded.
 
         Returns:
-            `str`: folder path of the repo snapshot.
+            `str` or list of [`DryRunFileInfo`]:
+                - If `dry_run=False`: Folder path of the repo snapshot.
+                - If `dry_run=True`: A list of [`DryRunFileInfo`] objects containing download information.
 
         Raises:
             [`~utils.RepositoryNotFoundError`]
@@ -5484,6 +5500,8 @@ def snapshot_download(
             ignore_patterns=ignore_patterns,
             max_workers=max_workers,
             tqdm_class=tqdm_class,
+            headers=self.headers,
+            dry_run=dry_run,
         )
 
     def get_safetensors_metadata(

tests/test_hf_api.py

@@ -3359,6 +3359,8 @@ def test_hf_hub_download_alias(self, mock: Mock) -> None:
             etag_timeout=10,
             local_files_only=False,
             headers=None,
+            tqdm_class=None,
+            dry_run=False,
         )
 
     @patch("huggingface_hub._snapshot_download.snapshot_download")
@@ -3385,6 +3387,8 @@ def test_snapshot_download_alias(self, mock: Mock) -> None:
             ignore_patterns=None,
             max_workers=8,
             tqdm_class=None,
+            headers=None,
+            dry_run=False,
         )