Skip to content

Commit

Permalink
utils: cleanup copyfile
Browse files Browse the repository at this point in the history
  • Loading branch information
skshetry committed Dec 11, 2023
1 parent b3a4b48 commit 7c9f6e5
Showing 1 changed file with 6 additions and 38 deletions.
44 changes: 6 additions & 38 deletions src/dvc_objects/fs/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -145,48 +145,12 @@ def makedirs(path, exist_ok: bool = False, mode: Optional[int] = None) -> None:
)


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,
callback: "Callback" = DEFAULT_CALLBACK,
) -> None:
"""Copy file with progress bar"""
total = os.stat(src).st_size

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

Expand All @@ -198,11 +162,15 @@ def copyfile(
except OSError:
pass

total = os.path.getsize(src)
if total < COPY_PBAR_MIN_SIZE:
shutil.copyfile(src, dest)
return

_copyfile_with_pbar(src, dest, total, **kwargs)
callback.set_size(total)
with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
wrapped = callback.wrap_attr(fdest, "write")
shutil.copyfileobj(fsrc, wrapped, length=LOCAL_CHUNK_SIZE) # type: ignore[misc]


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

0 comments on commit 7c9f6e5

Please sign in to comment.