Repository: huggingface/huggingface_hub
Pull request #1103 (2022-10-10 13:09:51)

respect NO_COLOR env var
Alternative to #1100. Makes `huggingface-cli` compliant with https://no-color.org.

HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/huggingface_hub/pr_1103). All of your documentation changes will be reflected on that endpoint.
osanseviero: Thanks for the contribution!

Patch:
diff --git a/src/huggingface_hub/__init__.py b/src/huggingface_hub/__init__.py
index 973b0219..f17bf3e3 100644
--- a/src/huggingface_hub/__init__.py
+++ b/src/huggingface_hub/__init__.py
@@ -164,6 +164,7 @@ _SUBMOD_ATTRS = {
         "Repository",
     ],
     "utils": [
+        "CacheNotFound",
         "CachedFileInfo",
         "CachedRepoInfo",
         "CachedRevisionInfo",
@@ -368,6 +369,7 @@ if TYPE_CHECKING:  # pragma: no cover
     from .utils import CachedFileInfo  # noqa: F401
     from .utils import CachedRepoInfo  # noqa: F401
     from .utils import CachedRevisionInfo  # noqa: F401
+    from .utils import CacheNotFound  # noqa: F401
     from .utils import CorruptedCacheException  # noqa: F401
     from .utils import DeleteCacheStrategy  # noqa: F401
     from .utils import HFCacheInfo  # noqa: F401
diff --git a/src/huggingface_hub/commands/_cli_utils.py b/src/huggingface_hub/commands/_cli_utils.py
index 18724d5a..bbf17e88 100644
--- a/src/huggingface_hub/commands/_cli_utils.py
+++ b/src/huggingface_hub/commands/_cli_utils.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Contains a utility for good-looking prints."""
+import os
 from typing import List, Union
 
 
@@ -39,6 +40,9 @@ class ANSI:
 
     @classmethod
     def _format(cls, s: str, code: str) -> str:
+        if os.environ.get("NO_COLOR"):
+            # See https://no-color.org/
+            return s
         return f"{code}{s}{cls._reset}"
 
 
diff --git a/src/huggingface_hub/commands/scan_cache.py b/src/huggingface_hub/commands/scan_cache.py
index edcbb198..fcf33337 100644
--- a/src/huggingface_hub/commands/scan_cache.py
+++ b/src/huggingface_hub/commands/scan_cache.py
@@ -24,7 +24,7 @@ import time
 from argparse import ArgumentParser
 from typing import Optional
 
-from ..utils import HFCacheInfo, scan_cache_dir
+from ..utils import CacheNotFound, HFCacheInfo, scan_cache_dir
 from . import BaseHuggingfaceCLICommand
 from ._cli_utils import ANSI, tabulate
 
@@ -59,9 +59,14 @@ class ScanCacheCommand(BaseHuggingfaceCLICommand):
         self.cache_dir: Optional[str] = args.dir
 
     def run(self):
-        t0 = time.time()
-        hf_cache_info = scan_cache_dir(self.cache_dir)
-        t1 = time.time()
+        try:
+            t0 = time.time()
+            hf_cache_info = scan_cache_dir(self.cache_dir)
+            t1 = time.time()
+        except CacheNotFound as exc:
+            cache_dir = exc.cache_dir
+            print(f"Cache directory not found: {cache_dir}")
+            return
 
         self._print_hf_cache_info_as_table(hf_cache_info)
 
diff --git a/src/huggingface_hub/utils/__init__.py b/src/huggingface_hub/utils/__init__.py
index 9f9d4e69..144b2402 100644
--- a/src/huggingface_hub/utils/__init__.py
+++ b/src/huggingface_hub/utils/__init__.py
@@ -20,6 +20,7 @@ from ._cache_manager import (
     CachedFileInfo,
     CachedRepoInfo,
     CachedRevisionInfo,
+    CacheNotFound,
     CorruptedCacheException,
     DeleteCacheStrategy,
     HFCacheInfo,
diff --git a/src/huggingface_hub/utils/_cache_manager.py b/src/huggingface_hub/utils/_cache_manager.py
index 43ca2d7a..d0fc72ee 100644
--- a/src/huggingface_hub/utils/_cache_manager.py
+++ b/src/huggingface_hub/utils/_cache_manager.py
@@ -31,6 +31,16 @@ logger = logging.get_logger(__name__)
 REPO_TYPE_T = Literal["model", "dataset", "space"]
 
 
+class CacheNotFound(Exception):
+    """Exception thrown when the Huggingface cache is not found."""
+
+    cache_dir = Union[str, Path]
+
+    def __init__(self, msg: str, cache_dir: Union[str, Path], *args, **kwargs):
+        super().__init__(msg, *args, **kwargs)
+        self.cache_dir = cache_dir
+
+
 class CorruptedCacheException(Exception):
     """Exception for any unexpected structure in the Huggingface cache-system."""
 
@@ -565,8 +575,11 @@ def scan_cache_dir(cache_dir: Optional[Union[str, Path]] = None) -> HFCacheInfo:
 
     Raises:
 
+        `CacheNotFound`
+          If the cache directory does not exist.
+
         [`ValueError`](https://docs.python.org/3/library/exceptions.html#ValueError)
-          If directory to scan is missing or is a file.
+          If the cache directory is a file, instead of a directory.
 
     </Tip>
 
@@ -577,9 +590,10 @@ def scan_cache_dir(cache_dir: Optional[Union[str, Path]] = None) -> HFCacheInfo:
 
     cache_dir = Path(cache_dir).expanduser().resolve()
     if not cache_dir.exists():
-        raise ValueError(
+        raise CacheNotFound(
             f"Cache directory not found: {cache_dir}. Please use `cache_dir` argument"
-            " or set `HUGGINGFACE_HUB_CACHE` environment variable."
+            " or set `HUGGINGFACE_HUB_CACHE` environment variable.",
+            cache_dir=cache_dir,
         )
 
     if cache_dir.is_file():


Test patch:
diff --git a/tests/test_utils_cache.py b/tests/test_utils_cache.py
index 7a17bc07..88a0957a 100644
--- a/tests/test_utils_cache.py
+++ b/tests/test_utils_cache.py
@@ -1,5 +1,6 @@
 import os
 import shutil
+import tempfile
 import time
 import unittest
 from pathlib import Path
@@ -11,6 +12,7 @@ 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._cache_manager import (
+    CacheNotFound,
     _format_size,
     _format_timesince,
     _try_delete_path,
@@ -34,8 +36,10 @@ class TestMissingCacheUtils(unittest.TestCase):
     cache_dir: Path
 
     def test_cache_dir_is_missing(self) -> None:
-        """Directory to scan does not exist raises ValueError."""
-        self.assertRaises(ValueError, scan_cache_dir, self.cache_dir / "does_not_exist")
+        """Directory to scan does not exist raises CacheNotFound."""
+        self.assertRaises(
+            CacheNotFound, scan_cache_dir, self.cache_dir / "does_not_exist"
+        )
 
     def test_cache_dir_is_a_file(self) -> None:
         """Directory to scan is a file raises ValueError."""
@@ -226,6 +230,27 @@ class TestValidCacheUtils(unittest.TestCase):
             expected_output.replace("-", "").split(),
         )
 
+    def test_cli_scan_missing_cache(self) -> None:
+        """Test output from CLI scan cache when cache does not exist.
+
+        End-to-end test just to see if output is in expected format.
+        """
+        tmp_dir = tempfile.mkdtemp()
+        os.rmdir(tmp_dir)
+
+        args = Mock()
+        args.verbose = 0
+        args.dir = tmp_dir
+
+        with capture_output() as output:
+            ScanCacheCommand(args).run()
+
+        expected_output = f"""
+        Cache directory not found: {tmp_dir}
+        """
+
+        self.assertListEqual(output.getvalue().split(), expected_output.split())
+
 
 @pytest.mark.usefixtures("fx_cache_dir")
 class TestCorruptedCacheUtils(unittest.TestCase):
diff --git a/tests/test_utils_cli.py b/tests/test_utils_cli.py
index cf02148c..e1d68118 100644
--- a/tests/test_utils_cli.py
+++ b/tests/test_utils_cli.py
@@ -1,9 +1,12 @@
+import os
 import unittest
+from unittest import mock
 
 from huggingface_hub.commands._cli_utils import ANSI, tabulate
 
 
 class TestCLIUtils(unittest.TestCase):
+    @mock.patch.dict(os.environ, {}, clear=True)
     def test_ansi_utils(self) -> None:
         """Test `ANSI` works as expected."""
         self.assertEqual(
@@ -26,6 +29,30 @@ class TestCLIUtils(unittest.TestCase):
             "\x1b[90m\x1b[1mthis is bold and grey\x1b[0m\x1b[0m",
         )
 
+    @mock.patch.dict(os.environ, {"NO_COLOR": "1"}, clear=True)
+    def test_ansi_no_color(self) -> None:
+        """Test `ANSI` respects `NO_COLOR` env var."""
+
+        self.assertEqual(
+            ANSI.bold("this is bold"),
+            "this is bold",
+        )
+
+        self.assertEqual(
+            ANSI.gray("this is gray"),
+            "this is gray",
+        )
+
+        self.assertEqual(
+            ANSI.red("this is red"),
+            "this is red",
+        )
+
+        self.assertEqual(
+            ANSI.gray(ANSI.bold("this is bold and grey")),
+            "this is bold and grey",
+        )
+
     def test_tabulate_utility(self) -> None:
         """Test `tabulate` works as expected."""
         rows = [[1, 2, 3], ["a very long value", "foo", "bar"], ["", 123, 456]]