objects: use separate Tqdm implementations

Now dvc.objects and dvc.data uses the implementation from dvc.objects
for Tqdm progress bars. This is temporary. :)

dvc/data/_progress.py

@@ -1,4 +1,4 @@
-from dvc.progress import Tqdm
+from dvc.objects._tqdm import Tqdm
 
 
 class QueryingProgress(Tqdm):

dvc/data/checkout.py

@@ -2,7 +2,7 @@ import logging
 from itertools import chain
 from typing import TYPE_CHECKING, List, Optional
 
-from dvc.objects.fs._callback import FsspecCallback
+from dvc.objects.fs.callbacks import Callback
 from dvc.objects.fs.generic import test_links, transfer
 
 from .diff import ROOT
@@ -154,7 +154,7 @@ class Link:
 
         cache.makedirs(cache.fs.path.parent(to_path))
         try:
-            with FsspecCallback.as_tqdm_callback(
+            with Callback.as_tqdm_callback(
                 callback,
                 desc=cache.fs.path.name(from_path),
                 bytes=True,

dvc/data/db/reference.py

@@ -9,8 +9,8 @@ from dvc.objects.fs import Schemes
 from ..reference import ReferenceHashFile
 
 if TYPE_CHECKING:
-    from dvc.objects.fs._callback import FsspecCallback
     from dvc.objects.fs.base import AnyFSPath, FileSystem
+    from dvc.objects.fs.callbacks import Callback
     from dvc.objects.hash_info import HashInfo
 
 logger = logging.getLogger(__name__)
@@ -58,7 +58,7 @@ class ReferenceObjectDB(ObjectDB):
         to_info: "AnyFSPath",
         hash_info: "HashInfo",
         hardlink: bool = False,
-        callback: "FsspecCallback" = None,
+        callback: "Callback" = None,
     ):
         if hash_info.isdir:
             return super()._add_file(

dvc/data/stage.py

@@ -6,11 +6,11 @@ from concurrent.futures import ThreadPoolExecutor
 from functools import partial
 from typing import TYPE_CHECKING, Dict, Optional, Tuple
 
+from dvc.objects._tqdm import Tqdm
 from dvc.objects.file import HashFile
 from dvc.objects.hash import hash_file
 from dvc.objects.hash_info import HashInfo
 from dvc.objects.meta import Meta
-from dvc.progress import Tqdm
 
 from .db.reference import ReferenceObjectDB
 
@@ -40,7 +40,7 @@ _STAGING_MEMFS_PATH = "dvc-staging"
 
 
 def _upload_file(from_fs_path, fs, odb, upload_odb, callback=None):
-    from dvc.objects.fs._callback import FsspecCallback
+    from dvc.objects.fs.callbacks import Callback
     from dvc.objects.fs.utils import tmp_fname
     from dvc.objects.hash import HashStreamFile
 
@@ -49,7 +49,7 @@ def _upload_file(from_fs_path, fs, odb, upload_odb, callback=None):
     with fs.open(from_fs_path, mode="rb") as stream:
         stream = HashStreamFile(stream)
         size = fs.size(from_fs_path)
-        with FsspecCallback.as_tqdm_callback(
+        with Callback.as_tqdm_callback(
             callback,
             desc=fs_path.name(from_fs_path),
             bytes=True,

dvc/data/transfer.py

@@ -5,8 +5,8 @@ from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional
 
 from funcy import split
 
+from dvc.objects._tqdm import Tqdm
 from dvc.objects.executors import ThreadPoolExecutor
-from dvc.progress import Tqdm
 
 if TYPE_CHECKING:
     from dvc.objects.db import ObjectDB

dvc/fs/_callback.py

@@ -1,7 +0,0 @@
-# pylint: disable=unused-import
-from dvc.objects.fs._callback import (  # noqa: F401
-    DEFAULT_CALLBACK,
-    FsspecCallback,
-    RichCallback,
-    TqdmCallback,
-)

dvc/fs/callbacks.py

@@ -0,0 +1,79 @@
+# pylint: disable=unused-import
+from contextlib import ExitStack
+from typing import TYPE_CHECKING, Optional
+
+from funcy import cached_property
+
+from dvc.objects.fs.callbacks import (  # noqa: F401
+    DEFAULT_CALLBACK,
+    Callback,
+    TqdmCallback,
+)
+
+if TYPE_CHECKING:
+    from dvc.ui._rich_progress import RichTransferProgress
+
+
+class RichCallback(Callback):
+    def __init__(
+        self,
+        size: Optional[int] = None,
+        value: int = 0,
+        progress: "RichTransferProgress" = None,
+        desc: str = None,
+        bytes: bool = False,  # pylint: disable=redefined-builtin
+        unit: str = None,
+        disable: bool = False,
+    ) -> None:
+        self._progress = progress
+        self.disable = disable
+        self._task_kwargs = {
+            "description": desc or "",
+            "bytes": bytes,
+            "unit": unit,
+            "total": size or 0,
+            "visible": False,
+            "progress_type": None if bytes else "summary",
+        }
+        self._stack = ExitStack()
+        super().__init__(size=size, value=value)
+
+    @cached_property
+    def progress(self):
+        from dvc.ui import ui
+        from dvc.ui._rich_progress import RichTransferProgress
+
+        if self._progress is not None:
+            return self._progress
+
+        progress = RichTransferProgress(
+            transient=True,
+            disable=self.disable,
+            console=ui.error_console,
+        )
+        return self._stack.enter_context(progress)
+
+    @cached_property
+    def task(self):
+        return self.progress.add_task(**self._task_kwargs)
+
+    def __enter__(self):
+        return self
+
+    def close(self):
+        self.progress.clear_task(self.task)
+        self._stack.close()
+
+    def call(self, hook_name=None, **kwargs):
+        self.progress.update(
+            self.task,
+            completed=self.value,
+            total=self.size,
+            visible=not self.disable,
+        )
+
+    def branch(self, path_1, path_2, kwargs, child: Optional[Callback] = None):
+        child = child or RichCallback(
+            progress=self.progress, desc=path_1, bytes=True
+        )
+        return super().branch(path_1, path_2, kwargs, child=child)

dvc/fs/data.py

@@ -5,8 +5,8 @@ import typing
 from fsspec import AbstractFileSystem
 from funcy import cached_property, wrap_prop
 
-from dvc.objects.fs._callback import DEFAULT_CALLBACK
 from dvc.objects.fs.base import FileSystem
+from dvc.objects.fs.callbacks import DEFAULT_CALLBACK
 
 if typing.TYPE_CHECKING:
     from dvc.types import AnyPath

dvc/fs/dvc.py

@@ -10,8 +10,8 @@ from typing import TYPE_CHECKING, Callable, Optional, Tuple, Type, Union
 from fsspec.spec import AbstractFileSystem
 from funcy import cached_property, wrap_prop, wrap_with
 
-from dvc.objects.fs._callback import DEFAULT_CALLBACK
 from dvc.objects.fs.base import FileSystem
+from dvc.objects.fs.callbacks import DEFAULT_CALLBACK
 from dvc.objects.fs.path import Path
 
 from .data import DataFileSystem

dvc/objects/_tqdm.py

@@ -0,0 +1,161 @@
+"""Manages progress bars."""
+import logging
+import os
+import re
+import sys
+from threading import RLock
+
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+tqdm.set_lock(RLock())
+
+
+def env2bool(var, undefined=False):
+    """
+    undefined: return value if env var is unset
+    """
+    var = os.getenv(var, None)
+    if var is None:
+        return undefined
+    return bool(re.search("1|y|yes|true", var, flags=re.I))
+
+
+class Tqdm(tqdm):
+    """
+    maximum-compatibility tqdm-based progressbars
+    """
+
+    BAR_FMT_DEFAULT = (
+        "{percentage:3.0f}% {desc}|{bar}|"
+        "{postfix[info]}{n_fmt}/{total_fmt}"
+        " [{elapsed}<{remaining}, {rate_fmt:>11}]"
+    )
+    # nested bars should have fixed bar widths to align nicely
+    BAR_FMT_DEFAULT_NESTED = (
+        "{percentage:3.0f}%|{bar:10}|{desc:{ncols_desc}.{ncols_desc}}"
+        "{postfix[info]}{n_fmt}/{total_fmt}"
+        " [{elapsed}<{remaining}, {rate_fmt:>11}]"
+    )
+    BAR_FMT_NOTOTAL = (
+        "{desc}{bar:b}|{postfix[info]}{n_fmt} [{elapsed}, {rate_fmt:>11}]"
+    )
+    BYTES_DEFAULTS = {
+        "unit": "B",
+        "unit_scale": True,
+        "unit_divisor": 1024,
+        "miniters": 1,
+    }
+
+    def __init__(
+        self,
+        iterable=None,
+        disable=None,
+        level=logging.ERROR,
+        desc=None,
+        leave=False,
+        bar_format=None,
+        bytes=False,  # pylint: disable=redefined-builtin
+        file=None,
+        total=None,
+        postfix=None,
+        **kwargs,
+    ):
+        """
+        bytes   : shortcut for
+            `unit='B', unit_scale=True, unit_divisor=1024, miniters=1`
+        desc  : persists after `close()`
+        level  : effective logging level for determining `disable`;
+            used only if `disable` is unspecified
+        disable  : If (default: None) or False,
+            will be determined by logging level.
+            May be overridden to `True` due to non-TTY status.
+            Skip override by specifying env var `DVC_IGNORE_ISATTY`.
+        kwargs  : anything accepted by `tqdm.tqdm()`
+        """
+        kwargs = kwargs.copy()
+        if bytes:
+            kwargs = {**self.BYTES_DEFAULTS, **kwargs}
+        else:
+            kwargs.setdefault("unit_scale", total > 999 if total else True)
+        if file is None:
+            file = sys.stderr
+        # auto-disable based on `logger.level`
+        if not disable:
+            disable = logger.getEffectiveLevel() > level
+        # auto-disable based on TTY
+        if (
+            not disable
+            and not env2bool("DVC_IGNORE_ISATTY")
+            and hasattr(file, "isatty")
+        ):
+            disable = not file.isatty()
+        super().__init__(
+            iterable=iterable,
+            disable=disable,
+            leave=leave,
+            desc=desc,
+            bar_format="!",
+            lock_args=(False,),
+            total=total,
+            **kwargs,
+        )
+        self.postfix = postfix or {"info": ""}
+        if bar_format is None:
+            if self.__len__():
+                self.bar_format = (
+                    self.BAR_FMT_DEFAULT_NESTED
+                    if self.pos
+                    else self.BAR_FMT_DEFAULT
+                )
+            else:
+                self.bar_format = self.BAR_FMT_NOTOTAL
+        else:
+            self.bar_format = bar_format
+        self.refresh()
+
+    def update_to(self, current, total=None):
+        if total:
+            self.total = total
+        self.update(current - self.n)
+
+    def close(self):
+        self.postfix["info"] = ""
+        # remove ETA (either unknown or zero); remove completed bar
+        self.bar_format = self.bar_format.replace("<{remaining}", "").replace(
+            "|{bar:10}|", " "
+        )
+        super().close()
+
+    def wrap_fn(self, fn, callback=None):
+        """
+        Returns a wrapped `fn` which calls `callback()` on each call.
+        `callback` is `self.update` by default.
+        """
+        if callback is None:
+            callback = self.update
+
+        def wrapped(*args, **kwargs):
+            res = fn(*args, **kwargs)
+            callback()
+            return res
+
+        return wrapped
+
+    @property
+    def format_dict(self):
+        """inject `ncols_desc` to fill the display width (`ncols`)"""
+        d = super().format_dict
+        ncols = d["ncols"] or 80
+        # assumes `bar_format` has max one of ("ncols_desc" & "ncols_info")
+        ncols_left = (
+            ncols - len(self.format_meter(ncols_desc=1, ncols_info=1, **d)) + 1
+        )
+        ncols_left = max(ncols_left, 0)
+        if ncols_left:
+            d["ncols_desc"] = d["ncols_info"] = ncols_left
+        else:
+            # work-around for zero-width description
+            d["ncols_desc"] = d["ncols_info"] = 1
+            d["prefix"] = ""
+        return d

dvc/objects/db.py

@@ -12,8 +12,8 @@ from .file import HashFile
 if TYPE_CHECKING:
     from typing import Tuple
 
-    from .fs._callback import FsspecCallback
     from .fs.base import AnyFSPath, FileSystem
+    from .fs.callbacks import Callback
     from .hash_info import HashInfo
 
 logger = logging.getLogger(__name__)
@@ -92,10 +92,10 @@ class ObjectDB:
         to_info: "AnyFSPath",
         _hash_info: "HashInfo",
         hardlink: bool = False,
-        callback: "FsspecCallback" = None,
+        callback: "Callback" = None,
     ):
         from .fs import generic
-        from .fs._callback import FsspecCallback
+        from .fs.callbacks import Callback
 
         self.makedirs(self.fs.path.parent(to_info))
         return generic.transfer(
@@ -104,7 +104,7 @@ class ObjectDB:
             self.fs,
             to_info,
             hardlink=hardlink,
-            callback=FsspecCallback.as_callback(callback),
+            callback=Callback.as_callback(callback),
         )
 
     def add(
@@ -114,9 +114,9 @@ class ObjectDB:
         hash_info: "HashInfo",
         hardlink: bool = False,
         verify: Optional[bool] = None,
-        callback: "FsspecCallback" = None,
+        callback: "Callback" = None,
     ):
-        from .fs._callback import FsspecCallback
+        from .fs.callbacks import Callback
 
         if self.read_only:
             raise ObjectDBPermissionError("Cannot add to read-only ODB")
@@ -130,7 +130,7 @@ class ObjectDB:
             pass
 
         cache_fs_path = self.hash_to_path(hash_info.value)
-        with FsspecCallback.as_tqdm_callback(
+        with Callback.as_tqdm_callback(
             callback,
             desc=fs.path.name(fs_path),
             bytes=True,

dvc/objects/fs/base.py

@@ -20,7 +20,7 @@ from typing import (
 from funcy import cached_property
 
 from ..executors import ThreadPoolExecutor
-from ._callback import DEFAULT_CALLBACK, FsspecCallback
+from .callbacks import DEFAULT_CALLBACK, Callback
 from .errors import RemoteMissingDepsError
 
 if TYPE_CHECKING:
@@ -351,7 +351,7 @@ class FileSystem:
         self,
         from_file: Union[AnyFSPath, "BinaryIO"],
         to_info: AnyFSPath,
-        callback: FsspecCallback = DEFAULT_CALLBACK,
+        callback: Callback = DEFAULT_CALLBACK,
         size: int = None,
         **kwargs,
     ) -> None:
@@ -371,7 +371,7 @@ class FileSystem:
         self,
         from_info: AnyFSPath,
         to_info: AnyFSPath,
-        callback: FsspecCallback = DEFAULT_CALLBACK,
+        callback: Callback = DEFAULT_CALLBACK,
         **kwargs,
     ) -> None:
         self.fs.get_file(from_info, to_info, callback=callback, **kwargs)
@@ -415,7 +415,7 @@ class FileSystem:
         self,
         from_info: Union[AnyFSPath, List[AnyFSPath]],
         to_info: Union[AnyFSPath, List[AnyFSPath]],
-        callback: "FsspecCallback" = DEFAULT_CALLBACK,
+        callback: "Callback" = DEFAULT_CALLBACK,
         recursive: bool = False,  # pylint: disable=unused-argument
         batch_size: int = None,
     ):
@@ -443,7 +443,7 @@ class FileSystem:
         self,
         from_info: Union[AnyFSPath, List[AnyFSPath]],
         to_info: Union[AnyFSPath, List[AnyFSPath]],
-        callback: "FsspecCallback" = DEFAULT_CALLBACK,
+        callback: "Callback" = DEFAULT_CALLBACK,
         recursive: bool = False,  # pylint: disable=unused-argument
         batch_size: int = None,
     ) -> None:

dvc/objects/fs/callbacks.py

@@ -10,15 +10,14 @@ if TYPE_CHECKING:
 
     from typing_extensions import ParamSpec
 
-    from dvc.progress import Tqdm
-    from dvc.ui._rich_progress import RichTransferProgress
+    from .._tqdm import Tqdm
 
     _P = ParamSpec("_P")
     _R = TypeVar("_R")
 
 
-class FsspecCallback(fsspec.Callback):
-    """FsspecCallback usable as a context manager, and a few helper methods."""
+class Callback(fsspec.Callback):
+    """Callback usable as a context manager, and a few helper methods."""
 
     @overload
     def wrap_attr(self, fobj: "BinaryIO", method: str = "read") -> "BinaryIO":
@@ -79,8 +78,8 @@ class FsspecCallback(fsspec.Callback):
 
     @classmethod
     def as_callback(
-        cls, maybe_callback: Optional["FsspecCallback"] = None
-    ) -> "FsspecCallback":
+        cls, maybe_callback: Optional["Callback"] = None
+    ) -> "Callback":
         if maybe_callback is None:
             return DEFAULT_CALLBACK
         return maybe_callback
@@ -88,33 +87,27 @@ class FsspecCallback(fsspec.Callback):
     @classmethod
     def as_tqdm_callback(
         cls,
-        callback: Optional["FsspecCallback"] = None,
+        callback: Optional["Callback"] = None,
         **tqdm_kwargs: Any,
-    ) -> "FsspecCallback":
+    ) -> "Callback":
         return callback or TqdmCallback(**tqdm_kwargs)
 
-    @classmethod
-    def as_rich_callback(
-        cls, callback: Optional["FsspecCallback"] = None, **rich_kwargs
-    ):
-        return callback or RichCallback(**rich_kwargs)
-
     def branch(
         self,
         path_1: str,
         path_2: str,
         kwargs: Dict[str, Any],
-        child: "FsspecCallback" = None,
-    ) -> "FsspecCallback":
+        child: "Callback" = None,
+    ) -> "Callback":
         child = kwargs["callback"] = child or DEFAULT_CALLBACK
         return child
 
 
-class NoOpCallback(FsspecCallback, fsspec.callbacks.NoOpCallback):
+class NoOpCallback(Callback, fsspec.callbacks.NoOpCallback):
     pass
 
 
-class TqdmCallback(FsspecCallback):
+class TqdmCallback(Callback):
     def __init__(
         self,
         size: Optional[int] = None,
@@ -130,7 +123,7 @@ class TqdmCallback(FsspecCallback):
 
     @cached_property
     def progress_bar(self):
-        from dvc.progress import Tqdm
+        from .._tqdm import Tqdm
 
         progress_bar = (
             self._progress_bar
@@ -159,77 +152,10 @@ class TqdmCallback(FsspecCallback):
         path_1: str,
         path_2: str,
         kwargs,
-        child: Optional[FsspecCallback] = None,
+        child: Optional[Callback] = None,
     ):
         child = child or TqdmCallback(bytes=True, desc=path_1)
         return super().branch(path_1, path_2, kwargs, child=child)
 
 
-class RichCallback(FsspecCallback):
-    def __init__(
-        self,
-        size: Optional[int] = None,
-        value: int = 0,
-        progress: "RichTransferProgress" = None,
-        desc: str = None,
-        bytes: bool = False,  # pylint: disable=redefined-builtin
-        unit: str = None,
-        disable: bool = False,
-    ) -> None:
-        self._progress = progress
-        self.disable = disable
-        self._task_kwargs = {
-            "description": desc or "",
-            "bytes": bytes,
-            "unit": unit,
-            "total": size or 0,
-            "visible": False,
-            "progress_type": None if bytes else "summary",
-        }
-        self._stack = ExitStack()
-        super().__init__(size=size, value=value)
-
-    @cached_property
-    def progress(self):
-        from dvc.ui import ui
-        from dvc.ui._rich_progress import RichTransferProgress
-
-        if self._progress is not None:
-            return self._progress
-
-        progress = RichTransferProgress(
-            transient=True,
-            disable=self.disable,
-            console=ui.error_console,
-        )
-        return self._stack.enter_context(progress)
-
-    @cached_property
-    def task(self):
-        return self.progress.add_task(**self._task_kwargs)
-
-    def __enter__(self):
-        return self
-
-    def close(self):
-        self.progress.clear_task(self.task)
-        self._stack.close()
-
-    def call(self, hook_name=None, **kwargs):
-        self.progress.update(
-            self.task,
-            completed=self.value,
-            total=self.size,
-            visible=not self.disable,
-        )
-
-    def branch(
-        self, path_1, path_2, kwargs, child: Optional[FsspecCallback] = None
-    ):
-        child = child or RichCallback(
-            progress=self.progress, desc=path_1, bytes=True
-        )
-        return super().branch(path_1, path_2, kwargs, child=child)
-
-
 DEFAULT_CALLBACK = NoOpCallback()

dvc/objects/fs/generic.py

@@ -3,13 +3,13 @@ import logging
 import os
 from typing import TYPE_CHECKING, List, Optional
 
-from ._callback import DEFAULT_CALLBACK
+from .callbacks import DEFAULT_CALLBACK
 from .implementations.local import LocalFileSystem
 from .utils import as_atomic
 
 if TYPE_CHECKING:
-    from ._callback import FsspecCallback
     from .base import AnyFSPath, FileSystem
+    from .callbacks import Callback
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +33,7 @@ def copy(
     from_path: "AnyFSPath",
     to_fs: "FileSystem",
     to_path: "AnyFSPath",
-    callback: "FsspecCallback" = DEFAULT_CALLBACK,
+    callback: "Callback" = DEFAULT_CALLBACK,
 ) -> None:
     if isinstance(from_fs, LocalFileSystem):
         return to_fs.put_file(from_path, to_path, callback=callback)
@@ -55,7 +55,7 @@ def _try_links(
     from_path: "AnyFSPath",
     to_fs: "FileSystem",
     to_path: "AnyFSPath",
-    callback: "FsspecCallback" = DEFAULT_CALLBACK,
+    callback: "Callback" = DEFAULT_CALLBACK,
 ) -> None:
     error = None
     while links:
@@ -85,7 +85,7 @@ def transfer(
     to_path: "AnyFSPath",
     hardlink: bool = False,
     links: Optional[List["str"]] = None,
-    callback: "FsspecCallback" = DEFAULT_CALLBACK,
+    callback: "Callback" = DEFAULT_CALLBACK,
 ) -> None:
     try:
         assert not (hardlink and links)

dvc/objects/fs/implementations/http.py

@@ -4,8 +4,8 @@ from typing import BinaryIO, Union
 
 from funcy import cached_property, memoize, wrap_with
 
-from .._callback import DEFAULT_CALLBACK, FsspecCallback
 from ..base import AnyFSPath, FileSystem
+from ..callbacks import DEFAULT_CALLBACK, Callback
 from ..errors import ConfigError
 
 
@@ -137,7 +137,7 @@ class HTTPFileSystem(FileSystem):
         self,
         from_file: Union[AnyFSPath, BinaryIO],
         to_info: AnyFSPath,
-        callback: FsspecCallback = DEFAULT_CALLBACK,
+        callback: Callback = DEFAULT_CALLBACK,
         size: int = None,
         **kwargs,
     ) -> None:

dvc/objects/fs/implementations/ssh.py

@@ -4,8 +4,8 @@ import threading
 
 from funcy import cached_property, memoize, silent, wrap_prop, wrap_with
 
-from .._callback import DEFAULT_CALLBACK
 from ..base import FileSystem
+from ..callbacks import DEFAULT_CALLBACK
 from ..utils import as_atomic
 
 DEFAULT_PORT = 22

dvc/objects/fs/utils.py

@@ -10,8 +10,8 @@ from typing import TYPE_CHECKING, Iterator
 from . import system
 
 if TYPE_CHECKING:
-    from ._callback import FsspecCallback
     from .base import AnyFSPath, FileSystem
+    from .callbacks import Callback
 
 
 logger = logging.getLogger(__name__)
@@ -140,7 +140,7 @@ def makedirs(path, exist_ok: bool = False, mode: int = None) -> None:
 def copyfile(
     src: "AnyFSPath",
     dest: "AnyFSPath",
-    callback: "FsspecCallback" = None,
+    callback: "Callback" = None,
     no_progress_bar: bool = False,
     name: str = None,
 ) -> None:
@@ -157,10 +157,10 @@ def copyfile(
     try:
         system.reflink(src, dest)
     except OSError:
-        from ._callback import FsspecCallback
+        from .callbacks import Callback
 
         with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
-            with FsspecCallback.as_tqdm_callback(
+            with Callback.as_tqdm_callback(
                 callback,
                 size=total,
                 bytes=True,

dvc/objects/hash.py

@@ -3,7 +3,7 @@ import io
 import logging
 from typing import TYPE_CHECKING, Any, BinaryIO, Dict, Optional, Tuple
 
-from .fs._callback import DEFAULT_CALLBACK, FsspecCallback, TqdmCallback
+from .fs.callbacks import DEFAULT_CALLBACK, Callback, TqdmCallback
 from .fs.implementations.local import localfs
 from .fs.utils import is_exec
 from .hash_info import HashInfo
@@ -86,7 +86,7 @@ def fobj_md5(
 def file_md5(
     fname: "AnyFSPath",
     fs: "FileSystem" = localfs,
-    callback: "FsspecCallback" = DEFAULT_CALLBACK,
+    callback: "Callback" = DEFAULT_CALLBACK,
     text: Optional[bool] = None,
 ) -> str:
     size = fs.size(fname) or 0
@@ -113,7 +113,7 @@ def _hash_file(
     fs_path: "AnyFSPath",
     fs: "FileSystem",
     name: str,
-    callback: "FsspecCallback" = DEFAULT_CALLBACK,
+    callback: "Callback" = DEFAULT_CALLBACK,
 ) -> Tuple["str", Dict[str, Any]]:
     info = _adapt_info(fs.info(fs_path), fs.protocol)
 
@@ -141,7 +141,7 @@ class LargeFileHashingCallback(TqdmCallback):
         self._logged = False
 
     # TqdmCallback force renders progress bar on `set_size`.
-    set_size = FsspecCallback.set_size
+    set_size = Callback.set_size
 
     def call(self, hook_name=None, **kwargs):
         if self.size and self.size > self.LARGE_FILE_SIZE:
@@ -160,7 +160,7 @@ def hash_file(
     fs: "FileSystem",
     name: str,
     state: "StateBase" = None,
-    callback: "FsspecCallback" = None,
+    callback: "Callback" = None,
 ) -> Tuple["Meta", "HashInfo"]:
     if state:
         meta, hash_info = state.get(fs_path, fs)

dvc/output.py

@@ -726,9 +726,9 @@ class Output:
             raise DvcException(msg.format(self.fs_path, name))
 
     def download(self, to, jobs=None):
-        from dvc.fs._callback import FsspecCallback
+        from dvc.fs.callbacks import Callback
 
-        with FsspecCallback.as_tqdm_callback(
+        with Callback.as_tqdm_callback(
             desc=f"Downloading {self.fs.path.name(self.fs_path)}",
             unit="files",
         ) as cb:

dvc/repo/get.py

@@ -21,7 +21,7 @@ def get(url, path, out=None, rev=None, jobs=None):
 
     from dvc.dvcfile import is_valid_filename
     from dvc.external_repo import external_repo
-    from dvc.fs._callback import FsspecCallback
+    from dvc.fs.callbacks import Callback
 
     out = resolve_output(path, out)
 
@@ -60,7 +60,7 @@ def get(url, path, out=None, rev=None, jobs=None):
                 fs = repo.dvcfs
                 fs_path = fs.from_os_path(path)
 
-            with FsspecCallback.as_tqdm_callback(
+            with Callback.as_tqdm_callback(
                 desc=f"Downloading {fs.path.name(path)}",
                 unit="files",
             ) as cb:

dvc/stage/cache.py

@@ -213,7 +213,7 @@ class StageCache:
         cached_stage.checkout()
 
     def transfer(self, from_odb, to_odb):
-        from dvc.fs._callback import FsspecCallback
+        from dvc.fs.callbacks import Callback
 
         from_fs = from_odb.fs
         to_fs = from_odb.fs
@@ -236,7 +236,7 @@ class StageCache:
 
             src_name = from_fs.path.name(src)
             parent_name = from_fs.path.name(from_fs.path.parent(src))
-            with FsspecCallback.as_tqdm_callback(
+            with Callback.as_tqdm_callback(
                 desc=src_name,
                 bytes=True,
             ) as cb:

dvc/utils/fs.py

@@ -157,10 +157,10 @@ def copyfile(src, dest, callback=None, no_progress_bar=False, name=None):
     try:
         system.reflink(src, dest)
     except OSError:
-        from dvc.fs._callback import FsspecCallback
+        from dvc.fs.callbacks import Callback
 
         with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
-            with FsspecCallback.as_tqdm_callback(
+            with Callback.as_tqdm_callback(
                 callback,
                 size=total,
                 bytes=True,

tests/func/test_fs.py

@@ -6,7 +6,7 @@ from os.path import join
 import pytest
 
 from dvc.fs import LocalFileSystem, get_cloud_fs
-from dvc.fs._callback import DEFAULT_CALLBACK, FsspecCallback
+from dvc.fs.callbacks import DEFAULT_CALLBACK, Callback
 from dvc.repo import Repo
 
 
@@ -246,7 +246,7 @@ def test_upload_callback(tmp_dir, dvc, cloud):
     fs = cls(**config)
     expected_size = os.path.getsize(tmp_dir / "foo")
 
-    callback = FsspecCallback()
+    callback = Callback()
     fs.put_file(
         (tmp_dir / "foo").fs_path,
         (cloud / "foo").fs_path,
@@ -262,7 +262,7 @@ def test_get_dir_callback(tmp_dir, dvc, cloud):
     fs = cls(**config)
     cloud.gen({"dir": {"foo": "foo", "bar": "bar"}})
 
-    callback = FsspecCallback()
+    callback = Callback()
     fs.get(
         (cloud / "dir").fs_path, (tmp_dir / "dir").fs_path, callback=callback
     )
@@ -278,7 +278,7 @@ def test_callback_on_dvcfs(tmp_dir, dvc, scm, mocker):
 
     fs = dvc.dvcfs
 
-    callback = FsspecCallback()
+    callback = Callback()
     fs.get(
         "dir",
         (tmp_dir / "dir2").fs_path,
@@ -289,7 +289,7 @@ def test_callback_on_dvcfs(tmp_dir, dvc, scm, mocker):
     assert callback.size == 2
     assert callback.value == 2
 
-    callback = FsspecCallback()
+    callback = Callback()
     branch = mocker.spy(callback, "branch")
     fs.get(
         os.path.join("dir", "foo"),
@@ -308,7 +308,7 @@ def test_callback_on_dvcfs(tmp_dir, dvc, scm, mocker):
 
     branch.reset_mock()
 
-    callback = FsspecCallback()
+    callback = Callback()
     branch = mocker.spy(callback, "branch")
     fs.get(
         os.path.join("dir", "bar"),
@@ -332,9 +332,8 @@ def test_callback_on_dvcfs(tmp_dir, dvc, scm, mocker):
 @pytest.mark.parametrize(
     "callback_factory, kwargs",
     [
-        (FsspecCallback.as_callback, {}),
-        (FsspecCallback.as_tqdm_callback, {"desc": "test"}),
-        (FsspecCallback.as_rich_callback, {"desc": "test"}),
+        (Callback.as_callback, {}),
+        (Callback.as_tqdm_callback, {"desc": "test"}),
     ],
 )
 def test_callback_with_none(request, api, callback_factory, kwargs, mocker):