Repository: huggingface/huggingface_hub
Pull request #1414 (2023-03-29 10:34:06)

Fix: notebook_login() does not update UI on Databricks
Fixes [1384](https://github.com/huggingface/huggingface_hub/issues/1384)

The PR adjusts notebook_login() to use the ipywidgets Box output to print the feedback.
This enables the logic to work in all notebooks that are compatible with ipywidgets - including Databricks where `clear_output` does not work.


Examples for the code flow:
<img width="1208" alt="Screenshot 2023-03-29 at 12 08 35" src="https://user-images.githubusercontent.com/91623108/228507934-ade816bd-d134-4ebb-bb57-f7aeb2c955e5.png">

<img width="1210" alt="Screenshot 2023-03-29 at 12 08 48" src="https://user-images.githubusercontent.com/91623108/228507949-cf981d78-22a1-4051-b6c6-1a401cfbcfa6.png">

<img width="1209" alt="Screenshot 2023-03-29 at 12 09 00" src="https://user-images.githubusercontent.com/91623108/228507962-c7f3e50a-2cda-4c48-8056-9ce8acfe732f.png">

<img width="1210" alt="Screenshot 2023-03-29 at 12 16 49" src="https://user-images.githubusercontent.com/91623108/228507977-bd6f7a7b-c6b7-4086-9130-d3e9cee3a44b.png">




HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/huggingface_hub/pr_1414). All of your documentation changes will be reflected on that endpoint.
julien-c: LGTM too and quite happy that we improve our support of Databricks / Databricks notebooks 👍
fwetdb: Thank you @Wauplin and everyone else involved!
The adjustments look good to me and I just double-checked them on Databricks.
So, good to merge from my side.

Patch:
diff --git a/src/huggingface_hub/_login.py b/src/huggingface_hub/_login.py
index b7693b30..12228b32 100644
--- a/src/huggingface_hub/_login.py
+++ b/src/huggingface_hub/_login.py
@@ -22,6 +22,7 @@ from .commands.delete_cache import _ask_for_confirmation_no_tui
 from .hf_api import HfApi
 from .utils import (
     HfFolder,
+    capture_output,
     is_google_colab,
     is_notebook,
     list_credential_helpers,
@@ -180,7 +181,7 @@ def notebook_login() -> None:
     """
     try:
         import ipywidgets.widgets as widgets  # type: ignore
-        from IPython.display import clear_output, display  # type: ignore
+        from IPython.display import display  # type: ignore
     except ImportError:
         raise ImportError(
             "The `notebook_login` function can only be used in a notebook (Jupyter or"
@@ -211,8 +212,16 @@ def notebook_login() -> None:
         add_to_git_credential = git_checkbox_widget.value
         # Erase token and clear value to make sure it's not saved in the notebook.
         token_widget.value = ""
-        clear_output()
-        _login(token, add_to_git_credential=add_to_git_credential)
+        # Hide inputs
+        login_token_widget.children = [widgets.Label("Connecting...")]
+        try:
+            with capture_output() as captured:
+                _login(token, add_to_git_credential=add_to_git_credential)
+            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)
 
@@ -235,13 +244,13 @@ def _login(token: str, add_to_git_credential: bool) -> None:
             set_git_credential(token)
             print(
                 "Your token has been saved in your configured git credential helpers"
-                f" ({','.join(list_credential_helpers())})."
+                + f" ({','.join(list_credential_helpers())})."
             )
         else:
             print("Token has not been saved to git credential helper.")
 
     HfFolder.save_token(token)
-    print("Your token has been saved to", HfFolder.path_token)
+    print(f"Your token has been saved to {HfFolder.path_token}")
     print("Login successful")
 
 
diff --git a/src/huggingface_hub/utils/__init__.py b/src/huggingface_hub/utils/__init__.py
index bfdf881b..f3f545d2 100644
--- a/src/huggingface_hub/utils/__init__.py
+++ b/src/huggingface_hub/utils/__init__.py
@@ -70,7 +70,7 @@ from ._runtime import (
     is_tf_available,
     is_torch_available,
 )
-from ._subprocess import run_interactive_subprocess, run_subprocess
+from ._subprocess import capture_output, run_interactive_subprocess, run_subprocess
 from ._validators import (
     HFValidationError,
     smoothly_deprecate_use_auth_token,
diff --git a/src/huggingface_hub/utils/_subprocess.py b/src/huggingface_hub/utils/_subprocess.py
index c9de3e32..5ec79365 100644
--- a/src/huggingface_hub/utils/_subprocess.py
+++ b/src/huggingface_hub/utils/_subprocess.py
@@ -16,7 +16,9 @@
 """Contains utilities to easily handle subprocesses in `huggingface_hub`."""
 import os
 import subprocess
+import sys
 from contextlib import contextmanager
+from io import StringIO
 from pathlib import Path
 from typing import IO, Generator, List, Optional, Tuple, Union
 
@@ -26,6 +28,26 @@ from .logging import get_logger
 logger = get_logger(__name__)
 
 
+@contextmanager
+def capture_output() -> Generator[StringIO, None, None]:
+    """Capture output that is printed to terminal.
+
+    Taken from https://stackoverflow.com/a/34738440
+
+    Example:
+    ```py
+    >>> with capture_output() as output:
+    ...     print("hello world")
+    >>> assert output.getvalue() == "hello world\n"
+    ```
+    """
+    output = StringIO()
+    previous_output = sys.stdout
+    sys.stdout = output
+    yield output
+    sys.stdout = previous_output
+
+
 def run_subprocess(
     command: Union[str, List[str]],
     folder: Optional[Union[str, Path]] = None,


Test patch:
diff --git a/tests/test_command_delete_cache.py b/tests/test_command_delete_cache.py
index 201bc20f..f5d3d902 100644
--- a/tests/test_command_delete_cache.py
+++ b/tests/test_command_delete_cache.py
@@ -16,9 +16,9 @@ from huggingface_hub.commands.delete_cache import (
     _manual_review_no_tui,
     _read_manual_review_tmp_file,
 )
-from huggingface_hub.utils import SoftTemporaryDirectory
+from huggingface_hub.utils import SoftTemporaryDirectory, capture_output
 
-from .testing_utils import capture_output, handle_injection
+from .testing_utils import handle_injection
 
 
 class TestDeleteCacheHelpers(unittest.TestCase):
diff --git a/tests/test_utils_cache.py b/tests/test_utils_cache.py
index 9ab80a26..429d49c2 100644
--- a/tests/test_utils_cache.py
+++ b/tests/test_utils_cache.py
@@ -9,7 +9,7 @@ import pytest
 
 from huggingface_hub._snapshot_download import snapshot_download
 from huggingface_hub.commands.scan_cache import ScanCacheCommand
-from huggingface_hub.utils import DeleteCacheStrategy, HFCacheInfo, scan_cache_dir
+from huggingface_hub.utils import DeleteCacheStrategy, HFCacheInfo, capture_output, scan_cache_dir
 from huggingface_hub.utils._cache_manager import (
     CacheNotFound,
     _format_size,
@@ -19,7 +19,6 @@ from huggingface_hub.utils._cache_manager import (
 
 from .testing_constants import TOKEN
 from .testing_utils import (
-    capture_output,
     rmtree_with_retry,
     with_production_testing,
     xfail_on_windows,
diff --git a/tests/testing_utils.py b/tests/testing_utils.py
index 2314f603..e93884db 100644
--- a/tests/testing_utils.py
+++ b/tests/testing_utils.py
@@ -2,16 +2,14 @@ import inspect
 import os
 import shutil
 import stat
-import sys
 import time
 import unittest
 import uuid
 from contextlib import contextmanager
 from enum import Enum
 from functools import wraps
-from io import StringIO
 from pathlib import Path
-from typing import Callable, Generator, Optional, Type, TypeVar, Union
+from typing import Callable, Optional, Type, TypeVar, Union
 from unittest.mock import Mock, patch
 
 import pytest
@@ -321,30 +319,6 @@ def xfail_on_windows(reason: str, raises: Optional[Type[Exception]] = None):
     return _inner_decorator
 
 
-@contextmanager
-def capture_output() -> Generator[StringIO, None, None]:
-    """Capture output that is printed to console.
-
-    Especially useful to test CLI commands.
-
-    Taken from https://stackoverflow.com/a/34738440
-
-    Example:
-    ```py
-    class TestHelloWorld(unittest.TestCase):
-        def test_hello_world(self):
-            with capture_output() as output:
-                print("hello world")
-            self.assertEqual(output.getvalue(), "hello world\n")
-    ```
-    """
-    output = StringIO()
-    previous_output = sys.stdout
-    sys.stdout = output
-    yield output
-    sys.stdout = previous_output
-
-
 T = TypeVar("T")