From da1e0b57b8b5d89d8c4a983d20ce5e6905060fb1 Mon Sep 17 00:00:00 2001 From: petrchmelar Date: Fri, 31 Jan 2025 14:56:27 +0100 Subject: [PATCH] import: fix broken progressbar when importing folder (#10678) Co-authored-by: Petr Chmelar Co-authored-by: skshetry <18718008+skshetry@users.noreply.github.com> --- dvc/fs/dvc.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py index 6d79cb8f9..f1370de95 100644 --- a/dvc/fs/dvc.py +++ b/dvc/fs/dvc.py @@ -500,7 +500,7 @@ class _DVCFileSystem(AbstractFileSystem): **kwargs, ) - def _get( # noqa: C901, PLR0912 + def _get( # noqa: C901, PLR0912, PLR0915 self, rpath, lpath, @@ -578,7 +578,8 @@ class _DVCFileSystem(AbstractFileSystem): kw = kwargs if isinstance(fs, DataFileSystem): kw = kw | {"info": info} - return fs.get_file(src, dest, callback=callback, **kw) + with callback.branched(src, dest) as child: + fs.get_file(src, dest, callback=child, **kw) if batch_size == 1: ctx: AbstractContextManager = nullcontext() @@ -589,7 +590,7 @@ class _DVCFileSystem(AbstractFileSystem): with ctx: it = ((fs, f) for fs, files in _files.items() for f in files) - deque(map_fn(get_file, it), maxlen=0) + deque(callback.wrap(map_fn(get_file, it)), maxlen=0) return result def get_file(self, rpath, lpath, **kwargs):