def chunks():
with open(from_file, "rb") as fd:
with Tqdm.wrapattr(
fd,
"read",
total=None
if no_progress_bar else os.path.getsize(from_file),
leave=False,
desc=to_info.url if name is None else name,
disable=no_progress_bar,
) as fd_wrapped:
while True:
chunk = fd_wrapped.read(self.CHUNK_SIZE)
if not chunk:
break
yield chunk
def _download(
self, from_info, to_file, name=None, no_progress_bar=False, **pbar_args
):
total = self.getsize(self._with_bucket(from_info))
with self.open(from_info, "rb") as fobj:
with Tqdm.wrapattr(
fobj,
"read",
desc=name,
disable=no_progress_bar,
bytes=True,
total=total,
**pbar_args,
) as wrapped:
with open(to_file, "wb") as fdest:
shutil.copyfileobj(wrapped, fdest, length=fobj.blocksize)
def _download(self, from_info, to_file, name=None, no_progress_bar=False):
response = self.request("GET", from_info.url, stream=True)
if response.status_code != 200:
raise HTTPError(response.status_code, response.reason)
with open(to_file, "wb") as fd:
with Tqdm.wrapattr(
fd,
"write",
total=None
if no_progress_bar else self._content_length(response),
leave=False,
desc=from_info.url if name is None else name,
disable=no_progress_bar,
) as fd_wrapped:
for chunk in response.iter_content(chunk_size=self.CHUNK_SIZE):
fd_wrapped.write(chunk)
def _download(self,
from_info,
to_file,
name=None,
no_progress_bar=False,
**_kwargs):
blob_client = self.blob_service.get_blob_client(
from_info.bucket, from_info.path)
total = blob_client.get_blob_properties().size
stream = blob_client.download_blob()
with open(to_file, "wb") as fobj:
with Tqdm.wrapattr(fobj,
"write",
desc=name,
total=total,
disable=no_progress_bar) as wrapped:
stream.readinto(wrapped)
def _upload(self,
from_file,
to_info,
name=None,
no_progress_bar=False,
**_kwargs):
blob_client = self.blob_service.get_blob_client(
to_info.bucket, to_info.path)
total = os.path.getsize(from_file)
with open(from_file, "rb") as fobj:
with Tqdm.wrapattr(fobj,
"read",
desc=name,
total=total,
disable=no_progress_bar) as wrapped:
blob_client.upload_blob(wrapped, overwrite=True)
def _upload(
self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
):
with self.hdfs(to_info) as hdfs:
tmp_file = tmp_fname(to_info.path)
total = os.path.getsize(from_file)
with open(from_file, "rb") as fobj:
with Tqdm.wrapattr(
fobj,
"read",
desc=name,
total=total,
disable=no_progress_bar,
) as wrapped:
with hdfs.open_output_stream(tmp_file) as sobj:
sobj.write(wrapped.read())
hdfs.move(tmp_file, to_info.path)
def upload_fobj(
self, fobj, to_info, no_progress_bar=False, size=None, **pbar_args
):
if not hasattr(self, "_upload_fobj"):
raise RemoteActionNotImplemented("upload_fobj", self.scheme)
with Tqdm.wrapattr(
fobj,
"read",
disable=no_progress_bar,
bytes=True,
total=size,
**pbar_args,
) as wrapped:
self._upload_fobj( # pylint: disable=no-member
wrapped, to_info, size=size
)
def _download(self,
from_info,
to_file,
name=None,
no_progress_bar=False,
**_kwargs):
with self.hdfs(from_info) as hdfs:
total = hdfs.info(from_info.path)["size"]
with open(to_file, "wb+") as fobj:
with Tqdm.wrapattr(
fobj,
"write",
desc=name,
total=total,
disable=no_progress_bar,
) as wrapped:
hdfs.download(from_info.path, wrapped)
def _upload_to_bucket(
bucket,
from_file,
to_info,
chunk_size=None,
name=None,
no_progress_bar=False,
):
blob = bucket.blob(to_info.path, chunk_size=chunk_size)
with open(from_file, mode="rb") as fobj:
with Tqdm.wrapattr(
fobj,
"read",
desc=name or to_info.path,
total=os.path.getsize(from_file),
disable=no_progress_bar,
) as wrapped:
blob.upload_from_file(wrapped)
def _download(self,
from_info,
to_file,
name=None,
no_progress_bar=False,
**_kwargs):
with self.hdfs(from_info) as hdfs:
file_info = hdfs.get_file_info(from_info.path)
total = file_info.size
with open(to_file, "wb+") as fobj:
with Tqdm.wrapattr(
fobj,
"write",
desc=name,
total=total,
disable=no_progress_bar,
) as wrapped:
with hdfs.open_input_stream(from_info.path) as sobj:
wrapped.write(sobj.read())
def _download(self,
from_info,
to_file,
name=None,
no_progress_bar=False,
**kwargs):
import shutil
from dvc.progress import Tqdm
with open(to_file, "wb+") as to_fobj:
with Tqdm.wrapattr(
to_fobj,
"write",
desc=name,
disable=no_progress_bar,
) as wrapped:
with self.open(from_info, "rb", **kwargs) as from_fobj:
shutil.copyfileobj(from_fobj, wrapped)
def _upload(self,
from_file,
to_info,
name=None,
no_progress_bar=False,
**_kwargs):
with self.hdfs(to_info) as hdfs:
tmp_file = tmp_fname(to_info.path)
total = os.path.getsize(from_file)
with open(from_file, "rb") as fobj:
with Tqdm.wrapattr(
fobj,
"read",
desc=name,
total=total,
disable=no_progress_bar,
) as wrapped:
hdfs.upload(tmp_file, wrapped)
hdfs.rename(tmp_file, to_info.path)
def gdrive_upload_file(
self, args, no_progress_bar=True, from_file="", progress_name=""
):
item = self.drive.CreateFile(
{"title": args["title"], "parents": [{"id": args["parent_id"]}]}
)
with open(from_file, "rb") as fobj:
total = os.path.getsize(from_file)
with Tqdm.wrapattr(
fobj,
"read",
desc=progress_name,
total=total,
disable=no_progress_bar,
) as wrapped:
# PyDrive doesn't like content property setting for empty files
# https://github.com/gsuitedevs/PyDrive/issues/121
if total:
item.content = wrapped
item.Upload()
return item
def _upload(self,
from_file,
to_info,
name=None,
no_progress_bar=False,
**_kwargs):
# First try to create parent directories
self.makedirs(to_info.parent)
file_size = os.path.getsize(from_file)
with open(from_file, "rb") as fd:
progress_context = (
nullcontext(fd) if file_size == 0 else Tqdm.wrapattr(
fd,
"read",
total=None if no_progress_bar else file_size,
leave=False,
desc=to_info.url if name is None else name,
disable=no_progress_bar,
))
with progress_context as fd_wrapped:
self._client.upload_to(buff=fd_wrapped,
remote_path=to_info.path)
def _gdrive_upload_file(
self,
parent_id,
title,
no_progress_bar=False,
from_file="",
progress_name="",
):
item = self._drive.CreateFile(
{"title": title, "parents": [{"id": parent_id}]}
)
with open(from_file, "rb") as fobj:
total = os.path.getsize(from_file)
with Tqdm.wrapattr(
fobj,
"read",
desc=progress_name,
total=total,
disable=no_progress_bar,
) as wrapped:
item.content = wrapped
item.Upload()
return item
