Skip to content

Commit

Permalink
fs.get_file: wrap callback when passing to transfer (#596)
Browse files Browse the repository at this point in the history
  • Loading branch information
skshetry authored Jan 31, 2025
1 parent 02739b0 commit 3b519f2
Showing 1 changed file with 14 additions and 2 deletions.
16 changes: 14 additions & 2 deletions src/dvc_data/fs.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from typing import Any, BinaryIO, NamedTuple, Optional

from fsspec import AbstractFileSystem
from fsspec.callbacks import DEFAULT_CALLBACK
from fsspec.callbacks import DEFAULT_CALLBACK, NoOpCallback

if typing.TYPE_CHECKING:
from dvc_objects.fs.base import AnyFSPath, FileSystem
Expand All @@ -22,6 +22,13 @@
logger = logging.getLogger(__name__)


class _WrappedCallback(NoOpCallback):
# check `_get_file` for more details
def branched(self, path_1, path_2, **kwargs):
# NOTE: only safe for a single use
return self.kw.get("callback", DEFAULT_CALLBACK)


class FileInfo(NamedTuple):
typ: str
storage: "ObjectStorage"
Expand Down Expand Up @@ -220,7 +227,12 @@ def get_file(
path,
fs,
os.fspath(lpath),
callback=callback,
# `transfer` supports uploading multiple files, so it uses the
# passed callback to iterate for no. of files.
# So, we wrap the given callback in a `NoOpCallback` and return it
# in `branch` so that file copy callback gets properly called.
# This is safe for transferring a single file.
callback=_WrappedCallback(callback=callback),
links=copy.copy(storage.odb.cache_types),
)
return
Expand Down

0 comments on commit 3b519f2

Please sign in to comment.