diff --git a/src/huggingface_hub/commands/upload.py b/src/huggingface_hub/commands/upload.py index 2fa6e70445..68640297ba 100644 --- a/src/huggingface_hub/commands/upload.py +++ b/src/huggingface_hub/commands/upload.py @@ -51,7 +51,7 @@ from huggingface_hub import logging from huggingface_hub._commit_scheduler import CommitScheduler from huggingface_hub.commands import BaseHuggingfaceCLICommand -from huggingface_hub.hf_api import create_repo, upload_file, upload_folder +from huggingface_hub.hf_api import HfApi from huggingface_hub.utils import disable_progress_bars, enable_progress_bars @@ -134,7 +134,7 @@ def __init__(self, args: Namespace) -> None: self.commit_message: Optional[str] = args.commit_message self.commit_description: Optional[str] = args.commit_description self.create_pr: bool = args.create_pr - self.token: Optional[str] = args.token + self.api: HfApi = HfApi(token=args.token, library_name="huggingface-cli") self.quiet: bool = args.quiet # disable warnings and progress bars # Check `--every` is valid @@ -222,7 +222,7 @@ def _upload(self) -> str: path_in_repo=path_in_repo, private=self.private, every=self.every, - token=self.token, + hf_api=self.api, ) print(f"Scheduling commits every {self.every} minutes to {scheduler.repo_id}.") try: # Block main thread until KeyboardInterrupt @@ -235,19 +235,18 @@ def _upload(self) -> str: # Otherwise, create repo and proceed with the upload if not os.path.isfile(self.local_path) and not os.path.isdir(self.local_path): raise FileNotFoundError(f"No such file or directory: '{self.local_path}'.") - repo_id = create_repo( - repo_id=self.repo_id, repo_type=self.repo_type, exist_ok=True, private=self.private, token=self.token + repo_id = self.api.create_repo( + repo_id=self.repo_id, repo_type=self.repo_type, exist_ok=True, private=self.private ).repo_id # File-based upload if os.path.isfile(self.local_path): - return upload_file( + return self.api.upload_file( path_or_fileobj=self.local_path, path_in_repo=self.path_in_repo, repo_id=repo_id, repo_type=self.repo_type, revision=self.revision, - token=self.token, commit_message=self.commit_message, commit_description=self.commit_description, create_pr=self.create_pr, @@ -255,13 +254,12 @@ def _upload(self) -> str: # Folder-based upload else: - return upload_folder( + return self.api.upload_folder( folder_path=self.local_path, path_in_repo=self.path_in_repo, repo_id=repo_id, repo_type=self.repo_type, revision=self.revision, - token=self.token, commit_message=self.commit_message, commit_description=self.commit_description, create_pr=self.create_pr, diff --git a/tests/test_cli.py b/tests/test_cli.py index 4117eb86f0..3d85b5c9c4 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -85,7 +85,7 @@ def test_upload_basic(self) -> None: self.assertEqual(cmd.commit_description, None) self.assertEqual(cmd.create_pr, False) self.assertEqual(cmd.every, None) - self.assertEqual(cmd.token, None) + self.assertEqual(cmd.api.token, None) self.assertEqual(cmd.quiet, False) def test_upload_with_all_options(self) -> None: @@ -135,7 +135,7 @@ def test_upload_with_all_options(self) -> None: self.assertEqual(cmd.commit_description, "My commit description") self.assertEqual(cmd.create_pr, True) self.assertEqual(cmd.every, 5) - self.assertEqual(cmd.token, "my-token") + self.assertEqual(cmd.api.token, "my-token") self.assertEqual(cmd.quiet, True) def test_upload_implicit_local_path_when_folder_exists(self) -> None: @@ -211,8 +211,8 @@ def test_every_as_float(self) -> None: cmd = UploadCommand(self.parser.parse_args(["upload", DUMMY_MODEL_ID, ".", "--every", "0.5"])) self.assertEqual(cmd.every, 0.5) - @patch("huggingface_hub.commands.upload.upload_folder") - @patch("huggingface_hub.commands.upload.create_repo") + @patch("huggingface_hub.commands.upload.HfApi.upload_folder") + @patch("huggingface_hub.commands.upload.HfApi.create_repo") def test_upload_folder_mock(self, create_mock: Mock, upload_mock: Mock) -> None: with SoftTemporaryDirectory() as cache_dir: cmd = UploadCommand( @@ -222,16 +222,13 @@ def test_upload_folder_mock(self, create_mock: Mock, upload_mock: Mock) -> None: ) cmd.run() - create_mock.assert_called_once_with( - repo_id="my-model", repo_type="model", exist_ok=True, private=True, token=None - ) + create_mock.assert_called_once_with(repo_id="my-model", repo_type="model", exist_ok=True, private=True) upload_mock.assert_called_once_with( folder_path=cache_dir, path_in_repo=".", repo_id=create_mock.return_value.repo_id, repo_type="model", revision=None, - token=None, commit_message=None, commit_description=None, create_pr=False, @@ -240,8 +237,8 @@ def test_upload_folder_mock(self, create_mock: Mock, upload_mock: Mock) -> None: delete_patterns=["*.json"], ) - @patch("huggingface_hub.commands.upload.upload_file") - @patch("huggingface_hub.commands.upload.create_repo") + @patch("huggingface_hub.commands.upload.HfApi.upload_file") + @patch("huggingface_hub.commands.upload.HfApi.create_repo") def test_upload_file_mock(self, create_mock: Mock, upload_mock: Mock) -> None: with SoftTemporaryDirectory() as cache_dir: file_path = Path(cache_dir) / "file.txt" @@ -254,7 +251,7 @@ def test_upload_file_mock(self, create_mock: Mock, upload_mock: Mock) -> None: cmd.run() create_mock.assert_called_once_with( - repo_id="my-dataset", repo_type="dataset", exist_ok=True, private=False, token=None + repo_id="my-dataset", repo_type="dataset", exist_ok=True, private=False ) upload_mock.assert_called_once_with( path_or_fileobj=str(file_path), @@ -262,13 +259,12 @@ def test_upload_file_mock(self, create_mock: Mock, upload_mock: Mock) -> None: repo_id=create_mock.return_value.repo_id, repo_type="dataset", revision=None, - token=None, commit_message=None, commit_description=None, create_pr=True, ) - @patch("huggingface_hub.commands.upload.create_repo") + @patch("huggingface_hub.commands.upload.HfApi.create_repo") def test_upload_missing_path(self, create_mock: Mock) -> None: cmd = UploadCommand(self.parser.parse_args(["upload", "my-model", "/path/to/missing_file", "logs/file.txt"])) with self.assertRaises(FileNotFoundError):