Skip to content

Commit

Permalink
fs: copyfile: don't touch callback at all for <1G files
Browse files Browse the repository at this point in the history
This again wastes a lot of time, as `set_size` might trigger
a pbar refresh.

Related #229
Related iterative/dvc#9914
  • Loading branch information
efiop committed Dec 1, 2023
1 parent 4a68565 commit 060a985
Showing 1 changed file with 40 additions and 26 deletions.
66 changes: 40 additions & 26 deletions src/dvc_objects/fs/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,48 +142,62 @@ def makedirs(path, exist_ok: bool = False, mode: Optional[int] = None) -> None:
)


def copyfile(
def _copyfile_with_pbar(
src: "AnyFSPath",
dest: "AnyFSPath",
total: int,
callback: Optional["Callback"] = None,
no_progress_bar: bool = False,
name: Optional[str] = None,
):
from .callbacks import Callback

name = name if name else os.path.basename(dest)

if callback:
callback.set_size(total)

with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
with Callback.as_tqdm_callback(
callback,
size=total,
bytes=True,
disable=no_progress_bar,
desc=name,
) as cb:
wrapped = cb.wrap_attr(fdest, "write")
while True:
buf = fsrc.read(LOCAL_CHUNK_SIZE)
if not buf:
break
wrapped.write(buf)

if callback:
callback.absolute_update(total)


def copyfile(
src: "AnyFSPath",
dest: "AnyFSPath",
**kwargs,
) -> None:
"""Copy file with progress bar"""
name = name if name else os.path.basename(dest)
total = os.stat(src).st_size

if os.path.isdir(dest):
dest = os.path.join(dest, os.path.basename(src))

if callback:
callback.set_size(total)

try:
system.reflink(src, dest)
return
except OSError:
if total < COPY_PBAR_MIN_SIZE:
shutil.copyfile(src, dest)
else:
from .callbacks import Callback

with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
with Callback.as_tqdm_callback(
callback,
size=total,
bytes=True,
disable=no_progress_bar,
desc=name,
) as cb:
wrapped = cb.wrap_attr(fdest, "write")
while True:
buf = fsrc.read(LOCAL_CHUNK_SIZE)
if not buf:
break
wrapped.write(buf)
pass

if callback:
callback.absolute_update(total)
if total < COPY_PBAR_MIN_SIZE:
shutil.copyfile(src, dest)
return

_copyfile_with_pbar(src, dest, total, **kwargs)


def tmp_fname(fname: "AnyFSPath" = "") -> "AnyFSPath":
Expand Down

0 comments on commit 060a985

Please sign in to comment.