diff --git a/docs/source/en/_redirects.yml b/docs/source/en/_redirects.yml index cfc5183087..4d419b475d 100644 --- a/docs/source/en/_redirects.yml +++ b/docs/source/en/_redirects.yml @@ -9,9 +9,9 @@ how-to-upstream: guides/upload search-the-hub: guides/search guides/manage_spaces: guides/manage-spaces package_reference/inference_api: package_reference/inference_client - +package_reference/login: package_reference/authentication # Alias for hf-transfer description hf_transfer: package_reference/environment_variables#hfhubenablehftransfer # Rename webhooks_server to webhooks -guides/webhooks_server: guides/webhooks \ No newline at end of file +guides/webhooks_server: guides/webhooks diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 9bb6c087dc..b4256b2a90 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -1,4 +1,4 @@ -- title: "Get started" +- title: 'Get started' sections: - local: index title: Home @@ -6,7 +6,7 @@ title: Quickstart - local: installation title: Installation -- title: "How-to guides" +- title: 'How-to guides' sections: - local: guides/overview title: Overview @@ -40,16 +40,16 @@ title: Integrate a library - local: guides/webhooks title: Webhooks -- title: "Conceptual guides" +- title: 'Conceptual guides' sections: - local: concepts/git_vs_http title: Git vs HTTP paradigm -- title: "Reference" +- title: 'Reference' sections: - local: package_reference/overview title: Overview - - local: package_reference/login - title: Login and logout + - local: package_reference/authentication + title: Authentication - local: package_reference/environment_variables title: Environment variables - local: package_reference/repository diff --git a/docs/source/en/guides/cli.md b/docs/source/en/guides/cli.md index 4995bd5d35..9fb355ae9d 100644 --- a/docs/source/en/guides/cli.md +++ b/docs/source/en/guides/cli.md @@ -92,7 +92,7 @@ Once you have your token, run the following command in your terminal: >>> huggingface-cli login ``` -This command will prompt you for a token. Copy-paste yours and press *Enter*. Then you'll be asked if the token should also be saved as a git credential. Press *Enter* again (default to yes) if you plan to use `git` locally. Finally, it will call the Hub to check that your token is valid and save it locally. +This command will prompt you for a token. Copy-paste yours and press *Enter*. Then, you'll be asked if the token should also be saved as a git credential. Press *Enter* again (default to yes) if you plan to use `git` locally. Finally, it will call the Hub to check that your token is valid and save it locally. ``` _| _| _| _| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _|_|_|_| _|_| _|_|_| _|_|_|_| @@ -102,7 +102,7 @@ _| _| _| _| _| _| _| _| _| _| _|_| _| _| _| _| _| _|_| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _| _| _| _|_|_| _|_|_|_| To log in, `huggingface_hub` requires a token generated from https://huggingface.co/settings/tokens . -Token: +Enter your token (input will not be visible): Add token as git credential? (Y/n) Token is valid (permission: write). Your token has been saved in your configured git credential helpers (store). @@ -114,11 +114,13 @@ Alternatively, if you want to log-in without being prompted, you can pass the to ```bash # Or using an environment variable ->>> huggingface-cli login --token $HUGGINGFACE_TOKEN --add-to-git-credential +>>> huggingface-cli login --token $HF_TOKEN --add-to-git-credential Token is valid (permission: write). +The token `token_name` has been saved to /home/wauplin/.cache/huggingface/stored_tokens Your token has been saved in your configured git credential helpers (store). Your token has been saved to /home/wauplin/.cache/huggingface/token Login successful +The current active token is: `token_name` ``` For more details about authentication, check out [this section](../quick-start#authentication). @@ -137,7 +139,7 @@ If you are not logged in, an error message will be printed. ## huggingface-cli logout -This commands logs you out. In practice, it will delete the token saved on your machine. +This commands logs you out. In practice, it will delete all tokens stored on your machine. If you want to remove a specific token, you can specify the token name as an argument. This command will not log you out if you are logged in using the `HF_TOKEN` environment variable (see [reference](../package_reference/environment_variables#hftoken)). If that is the case, you must unset the environment variable in your machine configuration. @@ -431,7 +433,7 @@ https://huggingface.co/Wauplin/my-cool-model/tree/main ## huggingface-cli repo-files -If you want to delete files from a Hugging Face repository, use the `huggingface-cli repo-files` command. +If you want to delete files from a Hugging Face repository, use the `huggingface-cli repo-files` command. ### Delete files @@ -439,17 +441,17 @@ The `huggingface-cli repo-files delete` sub-command allows you to dele Delete a folder : ```bash ->>> huggingface-cli repo-files Wauplin/my-cool-model delete folder/ +>>> huggingface-cli repo-files Wauplin/my-cool-model delete folder/ Files correctly deleted from repo. Commit: https://huggingface.co/Wauplin/my-cool-mo... ``` -Delete multiple files: +Delete multiple files: ```bash >>> huggingface-cli repo-files Wauplin/my-cool-model delete file.txt folder/pytorch_model.bin Files correctly deleted from repo. Commit: https://huggingface.co/Wauplin/my-cool-mo... ``` -Use Unix-style wildcards to delete sets of files: +Use Unix-style wildcards to delete sets of files: ```bash >>> huggingface-cli repo-files Wauplin/my-cool-model delete "*.txt" "folder/*.bin" Files correctly deleted from repo. Commit: https://huggingface.co/Wauplin/my-cool-mo... @@ -460,7 +462,7 @@ Files correctly deleted from repo. Commit: https://huggingface.co/Wauplin/my-coo To delete files from a repo you must be authenticated and authorized. By default, the token saved locally (using `huggingface-cli login`) will be used. If you want to authenticate explicitly, use the `--token` option: ```bash ->>> huggingface-cli repo-files --token=hf_**** Wauplin/my-cool-model delete file.txt +>>> huggingface-cli repo-files --token=hf_**** Wauplin/my-cool-model delete file.txt ``` ## huggingface-cli scan-cache diff --git a/docs/source/en/package_reference/login.md b/docs/source/en/package_reference/authentication.md similarity index 60% rename from docs/source/en/package_reference/login.md rename to docs/source/en/package_reference/authentication.md index 59b998ea71..1577b5bc29 100644 --- a/docs/source/en/package_reference/login.md +++ b/docs/source/en/package_reference/authentication.md @@ -2,9 +2,9 @@ rendered properly in your Markdown viewer. --> -# Login and logout +# Authentication -The `huggingface_hub` library allows users to programmatically login and logout the machine to the Hub. +The `huggingface_hub` library allows users to programmatically manage authentication to the Hub. This includes logging in, logging out, switching between tokens, and listing available tokens. For more details about authentication, check out [this section](../quick-start#authentication). @@ -23,3 +23,11 @@ For more details about authentication, check out [this section](../quick-start#a ## logout [[autodoc]] logout + +## auth_switch + +[[autodoc]] auth_switch + +## auth_list + +[[autodoc]] auth_list diff --git a/docs/source/en/quick-start.md b/docs/source/en/quick-start.md index 7105b58702..dbbe0d2a91 100644 --- a/docs/source/en/quick-start.md +++ b/docs/source/en/quick-start.md @@ -93,6 +93,19 @@ Once logged in, all requests to the Hub - even methods that don't necessarily re +### Manage multiple tokens locally + +You can save multiple tokens on your machine by simply logging in with the [`login`] command with each token. If you need to switch between these tokens locally, you can use the [`auth switch`] command: + +```bash +huggingface-cli auth switch +``` + +This command will prompt you to select a token by its name from a list of saved tokens. Once selected, the chosen token becomes the _active_ token, and it will be used for all interactions with the Hub. + + +You can list all available access tokens on your machine with `huggingface-cli auth list`. + ### Environment variable The environment variable `HF_TOKEN` can also be used to authenticate yourself. This is especially useful in a Space where you can set `HF_TOKEN` as a [Space secret](https://huggingface.co/docs/hub/spaces-overview#managing-secrets). diff --git a/src/huggingface_hub/__init__.py b/src/huggingface_hub/__init__.py index 413fac30fd..0d9613a1a6 100644 --- a/src/huggingface_hub/__init__.py +++ b/src/huggingface_hub/__init__.py @@ -63,6 +63,8 @@ "InferenceEndpointType", ], "_login": [ + "auth_list", + "auth_switch", "interpreter_login", "login", "logout", @@ -577,6 +579,8 @@ def __dir__(): InferenceEndpointType, # noqa: F401 ) from ._login import ( + auth_list, # noqa: F401 + auth_switch, # noqa: F401 interpreter_login, # noqa: F401 login, # noqa: F401 logout, # noqa: F401 diff --git a/src/huggingface_hub/_login.py b/src/huggingface_hub/_login.py index dbd2e15e65..24af592acf 100644 --- a/src/huggingface_hub/_login.py +++ b/src/huggingface_hub/_login.py @@ -15,6 +15,7 @@ import os import subprocess +import warnings from functools import partial from getpass import getpass from pathlib import Path @@ -33,7 +34,15 @@ set_git_credential, unset_git_credential, ) -from .utils._token import _get_token_from_environment, _get_token_from_google_colab +from .utils._auth import ( + _get_token_by_name, + _get_token_from_environment, + _get_token_from_file, + _get_token_from_google_colab, + _save_stored_tokens, + _save_token, + get_stored_tokens, +) logger = logging.get_logger(__name__) @@ -115,24 +124,35 @@ def login( interpreter_login(new_session=new_session, write_permission=write_permission) -def logout() -> None: +def logout(token_name: Optional[str] = None) -> None: """Logout the machine from the Hub. Token is deleted from the machine and removed from git credential. + + Args: + token_name (`str`, *optional*): + Name of the access token to logout from. If `None`, will logout from all saved access tokens. + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError): + If the access token name is not found. """ - if get_token() is None: + if get_token() is None and not get_stored_tokens(): # No active token and no saved access tokens print("Not logged in!") return + if not token_name: + # Delete all saved access tokens and token + for file_path in (constants.HF_TOKEN_PATH, constants.HF_STORED_TOKENS_PATH): + try: + Path(file_path).unlink() + except FileNotFoundError: + pass + print("Successfully logged out from all access tokens.") + else: + _logout_from_token(token_name) + print(f"Successfully logged out from access token: {token_name}.") - # Delete token from git credentials unset_git_credential() - # Delete token file - try: - Path(constants.HF_TOKEN_PATH).unlink() - except FileNotFoundError: - pass - # Check if still logged in if _get_token_from_google_colab() is not None: raise EnvironmentError( @@ -145,7 +165,70 @@ def logout() -> None: "To log out, you must clear out both `HF_TOKEN` and `HUGGING_FACE_HUB_TOKEN` environment variables." ) - print("Successfully logged out.") + +def auth_switch(token_name: str, add_to_git_credential: bool = False) -> None: + """Switch to a different access token. + + Args: + token_name (`str`): + Name of the access token to switch to. + add_to_git_credential (`bool`, defaults to `False`): + If `True`, token will be set as git credential. If no git credential helper + is configured, a warning will be displayed to the user. If `token` is `None`, + the value of `add_to_git_credential` is ignored and will be prompted again + to the end user. + + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError): + If the access token name is not found. + """ + token = _get_token_by_name(token_name) + if not token: + raise ValueError(f"Access token {token_name} not found in {constants.HF_STORED_TOKENS_PATH}") + # Write token to HF_TOKEN_PATH + _set_active_token(token_name, add_to_git_credential) + print(f"The current active token is: {token_name}") + token_from_environment = _get_token_from_environment() + if token_from_environment is not None and token_from_environment != token: + warnings.warn( + "The environment variable `HF_TOKEN` is set and will override the access token you've just switched to." + ) + + +def auth_list() -> None: + """List all stored access tokens.""" + tokens = get_stored_tokens() + + if not tokens: + print("No access tokens found.") + return + # Find current token + current_token = get_token() + current_token_name = None + for token_name in tokens: + if tokens.get(token_name) == current_token: + current_token_name = token_name + # Print header + max_offset = max(len("token"), max(len(token) for token in tokens)) + 2 + print(f" {{:<{max_offset}}}| {{:<15}}".format("name", "token")) + print("-" * (max_offset + 2) + "|" + "-" * 15) + + # Print saved access tokens + for token_name in tokens: + token = tokens.get(token_name, "") + masked_token = f"{token[:3]}****{token[-4:]}" if token != "" else token + is_current = "*" if token == current_token else " " + + print(f"{is_current} {{:<{max_offset}}}| {{:<15}}".format(token_name, masked_token)) + + if _get_token_from_environment(): + print( + "\nNote: Environment variable `HF_TOKEN` is set and is the current active token independently from the stored tokens listed above." + ) + elif current_token_name is None: + print( + "\nNote: No active token is set and no environment variable `HF_TOKEN` is found. Use `huggingface-cli login` to log in." + ) ### @@ -191,7 +274,11 @@ def interpreter_login(new_session: bool = True, write_permission: bool = False) token = getpass("Enter your token (input will not be visible): ") add_to_git_credential = _ask_for_confirmation_no_tui("Add token as git credential?") - _login(token=token, add_to_git_credential=add_to_git_credential, write_permission=write_permission) + _login( + token=token, + add_to_git_credential=add_to_git_credential, + write_permission=write_permission, + ) ### @@ -296,22 +383,75 @@ def login_token_event(t, write_permission: bool = False): ### -def _login(token: str, add_to_git_credential: bool, write_permission: bool = False) -> None: - from .hf_api import get_token_permission # avoid circular import +def _login( + token: str, + add_to_git_credential: bool, + write_permission: bool = False, +) -> None: + from .hf_api import whoami # avoid circular import if token.startswith("api_org"): raise ValueError("You must use your personal account token, not an organization token.") - permission = get_token_permission(token) - if permission is None: - raise ValueError("Invalid token passed!") - elif write_permission and permission != "write": + token_info = whoami(token) + permission = token_info["auth"]["accessToken"]["role"] + if write_permission and permission != "write": raise ValueError( "Token is valid but is 'read-only' and a 'write' token is required.\nPlease provide a new token with" " correct permission." ) print(f"Token is valid (permission: {permission}).") + token_name = token_info["auth"]["accessToken"]["displayName"] + # Store token locally + _save_token(token=token, token_name=token_name) + # Set active token + _set_active_token(token_name=token_name, add_to_git_credential=add_to_git_credential) + print("Login successful.") + if _get_token_from_environment(): + print( + "Note: Environment variable`HF_TOKEN` is set and is the current active token independently from the token you've just configured." + ) + else: + print(f"The current active token is: `{token_name}`") + + +def _logout_from_token(token_name: str) -> None: + """Logout from a specific access token. + + Args: + token_name (`str`): + The name of the access token to logout from. + Raises: + [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError): + If the access token name is not found. + """ + stored_tokens = get_stored_tokens() + # If there is no access tokens saved or the access token name is not found, do nothing + if not stored_tokens or token_name not in stored_tokens: + return + + token = stored_tokens.pop(token_name) + _save_stored_tokens(stored_tokens) + + if token == _get_token_from_file(): + warnings.warn(f"Active token '{token_name}' has been deleted.") + Path(constants.HF_TOKEN_PATH).unlink(missing_ok=True) + + +def _set_active_token( + token_name: str, + add_to_git_credential: bool, +) -> None: + """Set the active access token. + + Args: + token_name (`str`): + The name of the token to set as active. + """ + token = _get_token_by_name(token_name) + if not token: + raise ValueError(f"Token {token_name} not found in {constants.HF_STORED_TOKENS_PATH}") if add_to_git_credential: if _is_git_credential_helper_configured(): set_git_credential(token) @@ -321,13 +461,11 @@ def _login(token: str, add_to_git_credential: bool, write_permission: bool = Fal ) else: print("Token has not been saved to git credential helper.") - - # Save token + # Write token to HF_TOKEN_PATH path = Path(constants.HF_TOKEN_PATH) path.parent.mkdir(parents=True, exist_ok=True) path.write_text(token) print(f"Your token has been saved to {constants.HF_TOKEN_PATH}") - print("Login successful") def _current_token_okay(write_permission: bool = False): diff --git a/src/huggingface_hub/commands/user.py b/src/huggingface_hub/commands/user.py index 8cde3ac04c..c16bbefdb2 100644 --- a/src/huggingface_hub/commands/user.py +++ b/src/huggingface_hub/commands/user.py @@ -13,30 +13,39 @@ # limitations under the License. import subprocess from argparse import _SubParsersAction +from typing import Optional from requests.exceptions import HTTPError from huggingface_hub.commands import BaseHuggingfaceCLICommand -from huggingface_hub.constants import ( - ENDPOINT, - REPO_TYPES, - REPO_TYPES_URL_PREFIXES, - SPACES_SDK_TYPES, -) +from huggingface_hub.constants import ENDPOINT, REPO_TYPES, REPO_TYPES_URL_PREFIXES, SPACES_SDK_TYPES from huggingface_hub.hf_api import HfApi from .._login import ( # noqa: F401 # for backward compatibility # noqa: F401 # for backward compatibility NOTEBOOK_LOGIN_PASSWORD_HTML, NOTEBOOK_LOGIN_TOKEN_HTML_END, NOTEBOOK_LOGIN_TOKEN_HTML_START, + auth_list, + auth_switch, login, logout, notebook_login, ) -from ..utils import get_token +from ..utils import get_stored_tokens, get_token, logging from ._cli_utils import ANSI +logger = logging.get_logger(__name__) + +try: + from InquirerPy import inquirer + from InquirerPy.base.control import Choice + + _inquirer_py_available = True +except ImportError: + _inquirer_py_available = False + + class UserCommands(BaseHuggingfaceCLICommand): @staticmethod def register_subcommand(parser: _SubParsersAction): @@ -54,9 +63,31 @@ def register_subcommand(parser: _SubParsersAction): login_parser.set_defaults(func=lambda args: LoginCommand(args)) whoami_parser = parser.add_parser("whoami", help="Find out which huggingface.co account you are logged in as.") whoami_parser.set_defaults(func=lambda args: WhoamiCommand(args)) + logout_parser = parser.add_parser("logout", help="Log out") + logout_parser.add_argument( + "--token-name", + type=str, + help="Optional: Name of the access token to log out from.", + ) logout_parser.set_defaults(func=lambda args: LogoutCommand(args)) + auth_parser = parser.add_parser("auth", help="Other authentication related commands") + auth_subparsers = auth_parser.add_subparsers(help="Authentication subcommands") + auth_switch_parser = auth_subparsers.add_parser("switch", help="Switch between access tokens") + auth_switch_parser.add_argument( + "--token-name", + type=str, + help="Optional: Name of the access token to switch to.", + ) + auth_switch_parser.add_argument( + "--add-to-git-credential", + action="store_true", + help="Optional: Save token to git credential helper.", + ) + auth_switch_parser.set_defaults(func=lambda args: AuthSwitchCommand(args)) + auth_list_parser = auth_subparsers.add_parser("list", help="List all stored access tokens") + auth_list_parser.set_defaults(func=lambda args: AuthListCommand(args)) # new system: git-based repo system repo_parser = parser.add_parser("repo", help="{create} Commands to interact with your huggingface.co repos.") repo_subparsers = repo_parser.add_subparsers(help="huggingface.co repos related commands") @@ -95,12 +126,70 @@ def __init__(self, args): class LoginCommand(BaseUserCommand): def run(self): - login(token=self.args.token, add_to_git_credential=self.args.add_to_git_credential) + login( + token=self.args.token, + add_to_git_credential=self.args.add_to_git_credential, + ) class LogoutCommand(BaseUserCommand): def run(self): - logout() + logout(token_name=self.args.token_name) + + +class AuthSwitchCommand(BaseUserCommand): + def run(self): + token_name = self.args.token_name + if token_name is None: + token_name = self._select_token_name() + + if token_name is None: + print("No token name provided. Aborting.") + exit() + auth_switch(token_name, add_to_git_credential=self.args.add_to_git_credential) + + def _select_token_name(self) -> Optional[str]: + token_names = list(get_stored_tokens().keys()) + + if not token_names: + logger.error("No stored tokens found. Please login first.") + return None + + if _inquirer_py_available: + return self._select_token_name_tui(token_names) + # if inquirer is not available, use a simpler terminal UI + print("Available stored tokens:") + for i, token_name in enumerate(token_names, 1): + print(f"{i}. {token_name}") + while True: + try: + choice = input("Enter the number of the token to switch to (or 'q' to quit): ") + if choice.lower() == "q": + return None + index = int(choice) - 1 + if 0 <= index < len(token_names): + return token_names[index] + else: + print("Invalid selection. Please try again.") + except ValueError: + print("Invalid input. Please enter a number or 'q' to quit.") + + def _select_token_name_tui(self, token_names: list[str]) -> Optional[str]: + choices = [Choice(token_name, name=token_name) for token_name in token_names] + try: + return inquirer.select( + message="Select a token to switch to:", + choices=choices, + default=None, + ).execute() + except KeyboardInterrupt: + logger.info("Token selection cancelled.") + return None + + +class AuthListCommand(BaseUserCommand): + def run(self): + auth_list() class WhoamiCommand(BaseUserCommand): diff --git a/src/huggingface_hub/constants.py b/src/huggingface_hub/constants.py index 53a6699cef..b123517eeb 100644 --- a/src/huggingface_hub/constants.py +++ b/src/huggingface_hub/constants.py @@ -140,7 +140,7 @@ def _as_int(value: Optional[str]) -> Optional[int]: # See https://github.com/huggingface/huggingface_hub/issues/1232 _OLD_HF_TOKEN_PATH = os.path.expanduser("~/.huggingface/token") HF_TOKEN_PATH = os.environ.get("HF_TOKEN_PATH", os.path.join(HF_HOME, "token")) - +HF_STORED_TOKENS_PATH = os.path.join(os.path.dirname(HF_TOKEN_PATH), "stored_tokens") if _staging_mode: # In staging mode, we use a different cache to ensure we don't mix up production and staging data or tokens diff --git a/src/huggingface_hub/utils/__init__.py b/src/huggingface_hub/utils/__init__.py index 4efea4e253..8ad45734f1 100644 --- a/src/huggingface_hub/utils/__init__.py +++ b/src/huggingface_hub/utils/__init__.py @@ -35,6 +35,7 @@ ) from . import tqdm as _tqdm # _tqdm is the module +from ._auth import get_stored_tokens, get_token from ._cache_assets import cached_assets_path from ._cache_manager import ( CachedFileInfo, @@ -103,24 +104,9 @@ is_tf_available, is_torch_available, ) -from ._safetensors import ( - SafetensorsFileMetadata, - SafetensorsRepoMetadata, - TensorInfo, -) +from ._safetensors import SafetensorsFileMetadata, SafetensorsRepoMetadata, TensorInfo from ._subprocess import capture_output, run_interactive_subprocess, run_subprocess from ._telemetry import send_telemetry -from ._token import get_token from ._typing import is_jsonable, is_simple_optional_type, unwrap_simple_optional_type -from ._validators import ( - smoothly_deprecate_use_auth_token, - validate_hf_hub_args, - validate_repo_id, -) -from .tqdm import ( - are_progress_bars_disabled, - disable_progress_bars, - enable_progress_bars, - tqdm, - tqdm_stream_file, -) +from ._validators import smoothly_deprecate_use_auth_token, validate_hf_hub_args, validate_repo_id +from .tqdm import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars, tqdm, tqdm_stream_file diff --git a/src/huggingface_hub/utils/_token.py b/src/huggingface_hub/utils/_auth.py similarity index 66% rename from src/huggingface_hub/utils/_token.py rename to src/huggingface_hub/utils/_auth.py index 1faae9bc98..efdbd5c837 100644 --- a/src/huggingface_hub/utils/_token.py +++ b/src/huggingface_hub/utils/_auth.py @@ -13,11 +13,13 @@ # limitations under the License. """Contains an helper to get the token from machine (env variable, secret or config file).""" +import configparser +import logging import os import warnings from pathlib import Path from threading import Lock -from typing import Optional +from typing import Dict, Optional from .. import constants from ._runtime import is_colab_enterprise, is_google_colab @@ -27,6 +29,8 @@ _GOOGLE_COLAB_SECRET_LOCK = Lock() _GOOGLE_COLAB_SECRET: Optional[str] = None +logger = logging.getLogger(__name__) + def get_token() -> Optional[str]: """ @@ -68,8 +72,8 @@ def _get_token_from_google_colab() -> Optional[str]: return _GOOGLE_COLAB_SECRET try: - from google.colab import userdata - from google.colab.errors import Error as ColabError + from google.colab import userdata # type: ignore + from google.colab.errors import Error as ColabError # type: ignore except ImportError: return None @@ -121,6 +125,85 @@ def _get_token_from_file() -> Optional[str]: return None +def get_stored_tokens() -> Dict[str, str]: + """ + Returns the parsed INI file containing the access tokens. + The file is located at `HF_STORED_TOKENS_PATH`, defaulting to `~/.cache/huggingface/stored_tokens`. + If the file does not exist, an empty dictionary is returned. + + Returns: `Dict[str, str]` + Key is the token name and value is the token. + """ + tokens_path = Path(constants.HF_STORED_TOKENS_PATH) + if not tokens_path.exists(): + stored_tokens = {} + config = configparser.ConfigParser() + try: + config.read(tokens_path) + stored_tokens = {token_name: config.get(token_name, "hf_token") for token_name in config.sections()} + except configparser.Error as e: + logger.error(f"Error parsing stored tokens file: {e}") + stored_tokens = {} + return stored_tokens + + +def _save_stored_tokens(stored_tokens: Dict[str, str]) -> None: + """ + Saves the given configuration to the stored tokens file. + + Args: + stored_tokens (`Dict[str, str]`): + The stored tokens to save. Key is the token name and value is the token. + """ + stored_tokens_path = Path(constants.HF_STORED_TOKENS_PATH) + + # Write the stored tokens into an INI file + config = configparser.ConfigParser() + for token_name in sorted(stored_tokens.keys()): + config.add_section(token_name) + config.set(token_name, "hf_token", stored_tokens[token_name]) + + stored_tokens_path.parent.mkdir(parents=True, exist_ok=True) + with stored_tokens_path.open("w") as config_file: + config.write(config_file) + + +def _get_token_by_name(token_name: str) -> Optional[str]: + """ + Get the token by name. + + Args: + token_name (`str`): + The name of the token to get. + + Returns: + `str` or `None`: The token, `None` if it doesn't exist. + + """ + stored_tokens = get_stored_tokens() + if token_name not in stored_tokens: + return None + return _clean_token(stored_tokens[token_name]) + + +def _save_token(token: str, token_name: str) -> None: + """ + Save the given token. + + If the stored tokens file does not exist, it will be created. + Args: + token (`str`): + The token to save. + token_name (`str`): + The name of the token. + """ + tokens_path = Path(constants.HF_STORED_TOKENS_PATH) + stored_tokens = get_stored_tokens() + stored_tokens[token_name] = token + _save_stored_tokens(stored_tokens) + print(f"The token `{token_name}` has been saved to {tokens_path}") + + def _clean_token(token: Optional[str]) -> Optional[str]: """Clean token by removing trailing and leading spaces and newlines. diff --git a/src/huggingface_hub/utils/_headers.py b/src/huggingface_hub/utils/_headers.py index e76afb6cea..8b05e939db 100644 --- a/src/huggingface_hub/utils/_headers.py +++ b/src/huggingface_hub/utils/_headers.py @@ -19,6 +19,7 @@ from huggingface_hub.errors import LocalTokenNotFoundError from .. import constants +from ._auth import get_token from ._runtime import ( get_fastai_version, get_fastcore_version, @@ -31,7 +32,6 @@ is_tf_available, is_torch_available, ) -from ._token import get_token from ._validators import validate_hf_hub_args diff --git a/src/huggingface_hub/utils/_hf_folder.py b/src/huggingface_hub/utils/_hf_folder.py index 502b22658b..f4e4a98a6b 100644 --- a/src/huggingface_hub/utils/_hf_folder.py +++ b/src/huggingface_hub/utils/_hf_folder.py @@ -19,7 +19,7 @@ from typing import Optional from .. import constants -from ._token import get_token +from ._auth import get_token class HfFolder: diff --git a/src/huggingface_hub/utils/_runtime.py b/src/huggingface_hub/utils/_runtime.py index bdd19e6716..72b4dd4a2b 100644 --- a/src/huggingface_hub/utils/_runtime.py +++ b/src/huggingface_hub/utils/_runtime.py @@ -373,6 +373,7 @@ def dump_environment_info() -> Dict[str, Any]: info["HF_HUB_CACHE"] = constants.HF_HUB_CACHE info["HF_ASSETS_CACHE"] = constants.HF_ASSETS_CACHE info["HF_TOKEN_PATH"] = constants.HF_TOKEN_PATH + info["HF_STORED_TOKENS_PATH"] = constants.HF_STORED_TOKENS_PATH info["HF_HUB_OFFLINE"] = constants.HF_HUB_OFFLINE info["HF_HUB_DISABLE_TELEMETRY"] = constants.HF_HUB_DISABLE_TELEMETRY info["HF_HUB_DISABLE_PROGRESS_BARS"] = constants.HF_HUB_DISABLE_PROGRESS_BARS diff --git a/tests/test_auth.py b/tests/test_auth.py new file mode 100644 index 0000000000..fd1a18f641 --- /dev/null +++ b/tests/test_auth.py @@ -0,0 +1,162 @@ +import os +import tempfile +from unittest.mock import patch + +import pytest + +from huggingface_hub import constants +from huggingface_hub._login import _login, _set_active_token, auth_switch, logout +from huggingface_hub.utils._auth import _get_token_by_name, _get_token_from_file, _save_token, get_stored_tokens + +from .testing_constants import ENDPOINT_STAGING, OTHER_TOKEN, TOKEN + + +@pytest.fixture(autouse=True) +def use_tmp_file_paths(): + """ + Fixture to temporarily override HF_TOKEN_PATH, HF_STORED_TOKENS_PATH, and ENDPOINT. + + This fixture patches the constants in the huggingface_hub module to use the + specified paths and the staging endpoint. It also ensures that the files are + deleted after all tests in the module are completed. + """ + with tempfile.TemporaryDirectory() as tmp_hf_home: + hf_token_path = os.path.join(tmp_hf_home, "token") + hf_stored_tokens_path = os.path.join(tmp_hf_home, "stored_tokens") + with patch.multiple( + constants, + HF_TOKEN_PATH=hf_token_path, + HF_STORED_TOKENS_PATH=hf_stored_tokens_path, + ENDPOINT=ENDPOINT_STAGING, + ): + yield + + +class TestGetTokenByName: + def test_get_existing_token(self): + _save_token(TOKEN, "test_token") + token = _get_token_by_name("test_token") + assert token == TOKEN + + def test_get_non_existent_token(self): + assert _get_token_by_name("non_existent") is None + + +class TestSaveToken: + def test_save_new_token(self): + _save_token(TOKEN, "new_token") + + stored_tokens = get_stored_tokens() + assert "new_token" in stored_tokens + assert stored_tokens["new_token"] == TOKEN + + def test_overwrite_existing_token(self): + _save_token(TOKEN, "test_token") + _save_token("new_token", "test_token") + + assert _get_token_by_name("test_token") == "new_token" + + +class TestSetActiveToken: + def test_set_active_token_success(self): + _save_token(TOKEN, "test_token") + _set_active_token("test_token", add_to_git_credential=False) + assert _get_token_from_file() == TOKEN + + def test_set_active_token_non_existent(self): + non_existent_token = "non_existent" + with pytest.raises(ValueError, match="Token non_existent not found in .*"): + _set_active_token(non_existent_token, add_to_git_credential=False) + + +class TestLogin: + @patch( + "huggingface_hub.hf_api.whoami", + return_value={ + "auth": { + "accessToken": { + "displayName": "test_token", + "role": "write", + "createdAt": "2024-01-01T00:00:00.000Z", + } + } + }, + ) + def test_login_success(self, mock_whoami): + _login(TOKEN, add_to_git_credential=False) + + assert _get_token_by_name("test_token") == TOKEN + assert _get_token_from_file() == TOKEN + + @patch( + "huggingface_hub.hf_api.whoami", + return_value={ + "auth": { + "accessToken": { + "displayName": "test_token", + "role": "read", + "createdAt": "2024-01-01T00:00:00.000Z", + } + } + }, + ) + def test_login_errors(self, mock_whoami): + with pytest.raises(ValueError, match=r"Token is valid but is 'read-only' and a 'write' token is required.*"): + _login(TOKEN, add_to_git_credential=False, write_permission=True) + + +class TestLogout: + def test_logout_deletes_files(self): + _save_token(TOKEN, "test_token") + _set_active_token("test_token", add_to_git_credential=False) + + assert os.path.exists(constants.HF_TOKEN_PATH) + assert os.path.exists(constants.HF_STORED_TOKENS_PATH) + + logout() + # Check that both files are deleted + assert not os.path.exists(constants.HF_TOKEN_PATH) + assert not os.path.exists(constants.HF_STORED_TOKENS_PATH) + + def test_logout_specific_token(self): + # Create two tokens + _save_token(TOKEN, "token_1") + _save_token(OTHER_TOKEN, "token_2") + + logout("token_1") + # Check that HF_STORED_TOKENS_PATH still exists + assert os.path.exists(constants.HF_STORED_TOKENS_PATH) + # Check that token_1 is removed + stored_tokens = get_stored_tokens() + assert "token_1" not in stored_tokens + assert "token_2" in stored_tokens + + def test_logout_active_token(self): + _save_token(TOKEN, "active_token") + _set_active_token("active_token", add_to_git_credential=False) + + logout("active_token") + + # Check that both files are deleted + assert not os.path.exists(constants.HF_TOKEN_PATH) + stored_tokens = get_stored_tokens() + assert "active_token" not in stored_tokens + + +class TestAuthSwitch: + def test_auth_switch_existing_token(self): + # Add two access tokens + _save_token(TOKEN, "test_token_1") + _save_token(OTHER_TOKEN, "test_token_2") + # Set `test_token_1` as the active token + _set_active_token("test_token_1", add_to_git_credential=False) + + # Switch to `test_token_2` + auth_switch("test_token_2", add_to_git_credential=False) + + assert _get_token_from_file() == OTHER_TOKEN + + def test_auth_switch_nonexisting_token(self): + with patch("huggingface_hub.utils._auth._get_token_by_name", return_value=None): + with pytest.raises(ValueError): + auth_switch("nonexistent_token")