Skip to content

Commit

Permalink
expose and use name, no_progress_bar
Browse files Browse the repository at this point in the history
  • Loading branch information
casperdcl committed Nov 19, 2019
1 parent 5eefacd commit 28b99dc
Showing 1 changed file with 28 additions and 7 deletions.
35 changes: 28 additions & 7 deletions dvc/remote/gs.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,10 +48,20 @@ def wrapper(*args, **kwargs):


@dynamic_chunk_size
def _upload_to_bucket(bucket, from_file, to_info, chunk_size=None, **kwargs):
blob = bucket.blob(to_info.path, chunk_size=chunk_size, **kwargs)
def _upload_to_bucket(
bucket,
from_file,
to_info,
chunk_size=None,
name=None,
no_progress_bar=True,
):
blob = bucket.blob(to_info.path, chunk_size=chunk_size)
with Tqdm(
desc=to_info.path, total=os.path.getsize(from_file), bytes=True
desc=name or to_info.path,
total=os.path.getsize(from_file),
bytes=True,
disable=no_progress_bar,
) as pbar:
with io.open(from_file, mode="rb") as fd:
raw_read = fd.read
Expand Down Expand Up @@ -138,14 +148,25 @@ def exists(self, path_info):
paths = set(self._list_paths(path_info.bucket, path_info.path))
return any(path_info.path == path for path in paths)

def _upload(self, from_file, to_info, **_kwargs):
def _upload(self, from_file, to_info, name=None, no_progress_bar=True):
bucket = self.gs.bucket(to_info.bucket)
_upload_to_bucket(bucket, from_file, to_info)
_upload_to_bucket(
bucket,
from_file,
to_info,
name=name,
no_progress_bar=no_progress_bar,
)

def _download(self, from_info, to_file, **_kwargs):
def _download(self, from_info, to_file, name=None, no_progress_bar=True):
bucket = self.gs.bucket(from_info.bucket)
blob = bucket.get_blob(from_info.path)
with Tqdm(desc=from_info.path, total=blob.size, bytes=True) as pbar:
with Tqdm(
desc=name or from_info.path,
total=blob.size,
bytes=True,
disable=no_progress_bar,
) as pbar:
with io.open(to_file, mode="wb") as fd:
raw_write = fd.write

Expand Down

0 comments on commit 28b99dc

Please sign in to comment.