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 <lucainp@gmail.com>

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

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. </center>"""
 
 
-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,28 +251,34 @@ 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))
 
 
 ###
 # Login private helpers
 ###
 
 
-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.
 

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

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

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"