diff --git a/dvc/data/_progress.py b/dvc/data/_progress.py index 52202c270..e093e553f 100644 --- a/dvc/data/_progress.py +++ b/dvc/data/_progress.py @@ -1,4 +1,4 @@ -from dvc.progress import Tqdm +from dvc.objects._tqdm import Tqdm class QueryingProgress(Tqdm): diff --git a/dvc/data/checkout.py b/dvc/data/checkout.py index 0840e7d63..86597466a 100644 --- a/dvc/data/checkout.py +++ b/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, diff --git a/dvc/data/db/reference.py b/dvc/data/db/reference.py index a44f7bf24..d846925da 100644 --- a/dvc/data/db/reference.py +++ b/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( diff --git a/dvc/data/stage.py b/dvc/data/stage.py index f512357b0..4414f71b1 100644 --- a/dvc/data/stage.py +++ b/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, diff --git a/dvc/data/transfer.py b/dvc/data/transfer.py index 4b8d26268..9e2f970c4 100644 --- a/dvc/data/transfer.py +++ b/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 diff --git a/dvc/fs/_callback.py b/dvc/fs/_callback.py deleted file mode 100644 index a1189ca47..000000000 --- a/dvc/fs/_callback.py +++ /dev/null @@ -1,7 +0,0 @@ -# pylint: disable=unused-import -from dvc.objects.fs._callback import ( # noqa: F401 - DEFAULT_CALLBACK, - FsspecCallback, - RichCallback, - TqdmCallback, -) diff --git a/dvc/fs/callbacks.py b/dvc/fs/callbacks.py new file mode 100644 index 000000000..7a1a4aa2d --- /dev/null +++ b/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) diff --git a/dvc/fs/data.py b/dvc/fs/data.py index cba392eb2..8036619a9 100644 --- a/dvc/fs/data.py +++ b/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 diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py index b824a1a8b..a8d07700e 100644 --- a/dvc/fs/dvc.py +++ b/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 diff --git a/dvc/objects/_tqdm.py b/dvc/objects/_tqdm.py new file mode 100644 index 000000000..ad34a3cc1 --- /dev/null +++ b/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 diff --git a/dvc/objects/db.py b/dvc/objects/db.py index 7d35a092f..38d9771ad 100644 --- a/dvc/objects/db.py +++ b/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, diff --git a/dvc/objects/fs/base.py b/dvc/objects/fs/base.py index ff39b9e47..5451ef287 100644 --- a/dvc/objects/fs/base.py +++ b/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: diff --git a/dvc/objects/fs/_callback.py b/dvc/objects/fs/callbacks.py similarity index 58% rename from dvc/objects/fs/_callback.py rename to dvc/objects/fs/callbacks.py index 2e19aff66..86dc6383b 100644 --- a/dvc/objects/fs/_callback.py +++ b/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() diff --git a/dvc/objects/fs/generic.py b/dvc/objects/fs/generic.py index af4e89f9e..b76311ba6 100644 --- a/dvc/objects/fs/generic.py +++ b/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) diff --git a/dvc/objects/fs/implementations/http.py b/dvc/objects/fs/implementations/http.py index 0b6732dd8..27980f2d4 100644 --- a/dvc/objects/fs/implementations/http.py +++ b/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: diff --git a/dvc/objects/fs/implementations/ssh.py b/dvc/objects/fs/implementations/ssh.py index c186d67fd..8b93faf57 100644 --- a/dvc/objects/fs/implementations/ssh.py +++ b/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 diff --git a/dvc/objects/fs/utils.py b/dvc/objects/fs/utils.py index 37e0625eb..cc7d5fb6c 100644 --- a/dvc/objects/fs/utils.py +++ b/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, diff --git a/dvc/objects/hash.py b/dvc/objects/hash.py index b3cb02daf..59c30e4c5 100644 --- a/dvc/objects/hash.py +++ b/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) diff --git a/dvc/output.py b/dvc/output.py index 57254910e..193292a57 100644 --- a/dvc/output.py +++ b/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: diff --git a/dvc/repo/get.py b/dvc/repo/get.py index 9c0ece904..5f4029a2c 100644 --- a/dvc/repo/get.py +++ b/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: diff --git a/dvc/stage/cache.py b/dvc/stage/cache.py index 37bddc32b..dca761d38 100644 --- a/dvc/stage/cache.py +++ b/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: diff --git a/dvc/utils/fs.py b/dvc/utils/fs.py index 5e9e8aee2..ceeb77d27 100644 --- a/dvc/utils/fs.py +++ b/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, diff --git a/tests/func/test_fs.py b/tests/func/test_fs.py index 6df69466e..8dd399b83 100644 --- a/tests/func/test_fs.py +++ b/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):