From 0f0cc408a3e48074f330db36a1e99912c86332fc Mon Sep 17 00:00:00 2001 From: aliabid94 Date: Wed, 17 May 2023 06:49:37 -0700 Subject: [PATCH] Add new_session, write_permission args (#1476) * add new_session, write_permission args * code style * remove _is_valid_token * code quality --------- Co-authored-by: Lucain Pouget --- src/huggingface_hub/__init__.py | 2 + src/huggingface_hub/_login.py | 93 +++++++++++++++++++++------ src/huggingface_hub/hf_api.py | 22 ++++--- src/huggingface_hub/utils/__init__.py | 2 +- src/huggingface_hub/utils/_headers.py | 6 +- 5 files changed, 96 insertions(+), 29 deletions(-) diff --git a/src/huggingface_hub/__init__.py b/src/huggingface_hub/__init__.py index 600fc80cb7..55e4c3fa4b 100644 --- a/src/huggingface_hub/__init__.py +++ b/src/huggingface_hub/__init__.py @@ -158,6 +158,7 @@ "get_model_tags", "get_repo_discussions", "get_space_runtime", + "get_token_permission", "like", "list_datasets", "list_files_info", @@ -444,6 +445,7 @@ def __dir__(): get_model_tags, # noqa: F401 get_repo_discussions, # noqa: F401 get_space_runtime, # noqa: F401 + get_token_permission, # noqa: F401 like, # noqa: F401 list_datasets, # noqa: F401 list_files_info, # noqa: F401 diff --git a/src/huggingface_hub/_login.py b/src/huggingface_hub/_login.py index 12228b32ae..4bbb1c9235 100644 --- a/src/huggingface_hub/_login.py +++ b/src/huggingface_hub/_login.py @@ -14,12 +14,13 @@ """Contains methods to login to the Hub.""" import os import subprocess +from functools import partial from getpass import getpass from typing import Optional from .commands._cli_utils import ANSI from .commands.delete_cache import _ask_for_confirmation_no_tui -from .hf_api import HfApi +from .hf_api import get_token_permission from .utils import ( HfFolder, capture_output, @@ -36,7 +37,12 @@ logger = logging.get_logger(__name__) -def login(token: Optional[str] = None, add_to_git_credential: bool = False) -> None: +def login( + token: Optional[str] = None, + add_to_git_credential: bool = False, + new_session: bool = True, + write_permission: bool = False, +) -> None: """Login the machine to access the Hub. The `token` is persisted in cache and set as a git credential. Once done, the machine @@ -67,8 +73,10 @@ def login(token: Optional[str] = None, add_to_git_credential: bool = False) -> N 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. - - + new_session (`bool`, defaults to `True`): + If `True`, will request a token even if one is already saved on the machine. + write_permission (`bool`, defaults to `False`): + If `True`, requires a token with write permission. Raises: [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError) If an organization token is passed. Only personal account tokens are valid @@ -85,11 +93,11 @@ def login(token: Optional[str] = None, add_to_git_credential: bool = False) -> N " `add_to_git_credential=True` if you want to set the git" " credential as well." ) - _login(token, add_to_git_credential=add_to_git_credential) + _login(token, add_to_git_credential=add_to_git_credential, write_permission=write_permission) elif is_notebook(): - notebook_login() + notebook_login(new_session=new_session, write_permission=write_permission) else: - interpreter_login() + interpreter_login(new_session=new_session, write_permission=write_permission) def logout() -> None: @@ -111,7 +119,7 @@ def logout() -> None: ### -def interpreter_login() -> None: +def interpreter_login(new_session: bool = True, write_permission: bool = False) -> None: """ Displays a prompt to login to the HF website and store the token. @@ -120,7 +128,18 @@ def interpreter_login() -> None: instead of a notebook widget. For more details, see [`login`]. + + Args: + new_session (`bool`, defaults to `True`): + If `True`, will request a token even if one is already saved on the machine. + write_permission (`bool`, defaults to `False`): + If `True`, requires a token with write permission. + """ + if not new_session and _current_token_okay(write_permission=write_permission): + print("User is already logged in.") + return + print(""" _| _| _| _| _|_|_| _|_|_| _|_|_| _| _| _|_|_| _|_|_|_| _|_| _|_|_| _|_|_|_| _| _| _| _| _| _| _| _|_| _| _| _| _| _| _| _| @@ -142,7 +161,7 @@ def interpreter_login() -> None: token = getpass("Token: ") 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) + _login(token=token, add_to_git_credential=add_to_git_credential, write_permission=write_permission) ### @@ -169,7 +188,7 @@ def interpreter_login() -> None: notebooks. """ -def notebook_login() -> None: +def notebook_login(new_session: bool = True, write_permission: bool = False) -> None: """ Displays a widget to login to the HF website and store the token. @@ -178,6 +197,12 @@ def notebook_login() -> None: instead of a prompt in the terminal. For more details, see [`login`]. + + Args: + new_session (`bool`, defaults to `True`): + If `True`, will request a token even if one is already saved on the machine. + write_permission (`bool`, defaults to `False`): + If `True`, requires a token with write permission. """ try: import ipywidgets.widgets as widgets # type: ignore @@ -187,6 +212,9 @@ def notebook_login() -> None: "The `notebook_login` function can only be used in a notebook (Jupyter or" " Colab) and you need the `ipywidgets` module: `pip install ipywidgets`." ) + if not new_session and _current_token_okay(write_permission=write_permission): + print("User is already logged in.") + return box_layout = widgets.Layout(display="flex", flex_flow="column", align_items="center", width="50%") @@ -207,7 +235,14 @@ def notebook_login() -> None: display(login_token_widget) # On click events - def login_token_event(t): + def login_token_event(t, write_permission: bool = False): + """ + Event handler for the login button. + + Args: + write_permission (`bool`, defaults to `False`): + If `True`, requires a token with write permission. + """ token = token_widget.value add_to_git_credential = git_checkbox_widget.value # Erase token and clear value to make sure it's not saved in the notebook. @@ -216,14 +251,14 @@ def login_token_event(t): login_token_widget.children = [widgets.Label("Connecting...")] try: with capture_output() as captured: - _login(token, add_to_git_credential=add_to_git_credential) + _login(token, add_to_git_credential=add_to_git_credential, write_permission=write_permission) message = captured.getvalue() except Exception as error: message = str(error) # Print result (success message or error) login_token_widget.children = [widgets.Label(line) for line in message.split("\n") if line.strip()] - token_finish_button.on_click(login_token_event) + token_finish_button.on_click(partial(login_token_event, write_permission=write_permission)) ### @@ -231,13 +266,19 @@ def login_token_event(t): ### -def _login(token: str, add_to_git_credential: bool) -> None: - hf_api = HfApi() +def _login(token: str, add_to_git_credential: bool, write_permission: bool = False) -> None: if token.startswith("api_org"): - raise ValueError("You must use your personal account token.") - if not hf_api._is_valid_token(token=token): + 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!") - print("Token is valid.") + elif 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}).") if add_to_git_credential: if _is_git_credential_helper_configured(): @@ -254,6 +295,22 @@ def _login(token: str, add_to_git_credential: bool) -> None: print("Login successful") +def _current_token_okay(write_permission: bool = False): + """Check if the current token is valid. + + Args: + write_permission (`bool`, defaults to `False`): + If `True`, requires a token with write permission. + + Returns: + `bool`: `True` if the current token is valid, `False` otherwise. + """ + permission = get_token_permission() + if permission is None or (write_permission and permission != "write"): + return False + return True + + def _is_git_credential_helper_configured() -> bool: """Check if a git credential helper is configured. diff --git a/src/huggingface_hub/hf_api.py b/src/huggingface_hub/hf_api.py index dccad33b04..c37c6990e8 100644 --- a/src/huggingface_hub/hf_api.py +++ b/src/huggingface_hub/hf_api.py @@ -30,6 +30,7 @@ from huggingface_hub.utils import ( IGNORE_GIT_FOLDER_PATTERNS, EntryNotFoundError, + LocalTokenNotFoundError, RepositoryNotFoundError, experimental, get_session, @@ -855,22 +856,24 @@ def whoami(self, token: Optional[str] = None) -> Dict: ) from e return r.json() - def _is_valid_token(self, token: str) -> bool: + def get_token_permission(self, token: Optional[str] = None) -> Literal["read", "write", None]: """ - Determines whether `token` is a valid token or not. + Check if a given `token` is valid and return its permissions. + + For more details about tokens, please refer to https://huggingface.co/docs/hub/security-tokens#what-are-user-access-tokens. Args: - token (`str`): - The token to check for validity. + token (`str`, *optional*): + The token to check for validity. Defaults to the one saved locally. Returns: - `bool`: `True` if valid, `False` otherwise. + `Literal["read", "write", None]`: Permission granted by the token ("read" or "write"). Returns `None` if no + token passed or token is invalid. """ try: - self.whoami(token=token) - return True - except HTTPError: - return False + return self.whoami(token=token)["auth"]["accessToken"]["role"] + except (LocalTokenNotFoundError, HTTPError): + return None def get_model_tags(self) -> ModelTags: "Gets all valid model tags as a nested namespace object" @@ -4866,6 +4869,7 @@ def _parse_revision_from_pr_url(pr_url: str) -> str: api = HfApi() whoami = api.whoami +get_token_permission = api.get_token_permission list_models = api.list_models model_info = api.model_info diff --git a/src/huggingface_hub/utils/__init__.py b/src/huggingface_hub/utils/__init__.py index cec24731b5..5e43d975bf 100644 --- a/src/huggingface_hub/utils/__init__.py +++ b/src/huggingface_hub/utils/__init__.py @@ -41,7 +41,7 @@ ) from ._fixes import SoftTemporaryDirectory, yaml_dump from ._git_credential import list_credential_helpers, set_git_credential, unset_git_credential -from ._headers import build_hf_headers, get_token_to_send +from ._headers import build_hf_headers, get_token_to_send, LocalTokenNotFoundError from ._hf_folder import HfFolder from ._http import configure_http_backend, get_session, http_backoff from ._pagination import paginate diff --git a/src/huggingface_hub/utils/_headers.py b/src/huggingface_hub/utils/_headers.py index a89237bece..846cca3f1d 100644 --- a/src/huggingface_hub/utils/_headers.py +++ b/src/huggingface_hub/utils/_headers.py @@ -32,6 +32,10 @@ from ._validators import validate_hf_hub_args +class LocalTokenNotFoundError(EnvironmentError): + """Raised if local token is required but not found.""" + + @validate_hf_hub_args def build_hf_headers( *, @@ -146,7 +150,7 @@ def get_token_to_send(token: Optional[Union[bool, str]]) -> Optional[str]: # Case token is explicitly required if token is True: if cached_token is None: - raise EnvironmentError( + raise LocalTokenNotFoundError( "Token is required (`token=True`), but no token found. You" " need to provide a token or be logged in to Hugging Face with" " `huggingface-cli login` or `huggingface_hub.login`. See"