Merge pull request #2875 from casperdcl/stream-progress

Stream progress
iterative · Dec 3, 2019 · f142188 · f142188
2 parents f91aa06 + c622287
commit f142188
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 75 deletions.
diff --git a/dvc/remote/gdrive/__init__.py → dvc/remote/gdrive.py b/dvc/remote/gdrive/__init__.py → dvc/remote/gdrive.py
@@ -8,7 +8,7 @@
 from funcy import retry, compose, decorator, wrap_with
 from funcy.py3 import cat
 
-from dvc.remote.gdrive.utils import TrackFileReadProgress, FOLDER_MIME_TYPE
+from dvc.progress import Tqdm
 from dvc.scheme import Schemes
 from dvc.path_info import CloudURLInfo
 from dvc.remote.base import RemoteBASE
@@ -17,6 +17,7 @@
 from dvc.utils import tmp_fname
 
 logger = logging.getLogger(__name__)
+FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
 
 
 class GDriveRetriableError(DvcException):
@@ -92,30 +93,34 @@ def gdrive_upload_file(
         item = self.drive.CreateFile(
             {"title": args["title"], "parents": [{"id": args["parent_id"]}]}
         )
-        self.upload_file(item, no_progress_bar, from_file, progress_name)
-        return item
 
-    def upload_file(self, item, no_progress_bar, from_file, progress_name):
-        with open(from_file, "rb") as opened_file:
-            if not no_progress_bar:
-                opened_file = TrackFileReadProgress(progress_name, opened_file)
-            # PyDrive doesn't like content property setting for empty files
-            # https://github.com/gsuitedevs/PyDrive/issues/121
-            if os.stat(from_file).st_size:
-                item.content = opened_file
-            item.Upload()
+        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
 
     @gdrive_retry
     def gdrive_download_file(
         self, file_id, to_file, progress_name, no_progress_bar
     ):
-        from dvc.progress import Tqdm
-
         gdrive_file = self.drive.CreateFile({"id": file_id})
+        bar_format = (
+            "Donwloading {desc:{ncols_desc}.{ncols_desc}}... "
+            + Tqdm.format_sizeof(int(gdrive_file["fileSize"]), "B", 1024)
+        )
         with Tqdm(
-            desc=progress_name,
-            total=int(gdrive_file["fileSize"]),
-            disable=no_progress_bar,
+            bar_format=bar_format, desc=progress_name, disable=no_progress_bar
         ):
             gdrive_file.GetContentFile(to_file)
 

diff --git a/dvc/remote/gdrive/utils.py b/dvc/remote/gdrive/utils.py
diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py
@@ -58,23 +58,15 @@ def _upload_to_bucket(
     no_progress_bar=True,
 ):
     blob = bucket.blob(to_info.path, chunk_size=chunk_size)
-    with Tqdm(
-        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 fobj:
-            raw_read = fobj.read
-
-            def read(size=chunk_size):
-                res = raw_read(size)
-                if res:
-                    pbar.update(len(res))
-                return res
-
-            fobj.read = read
-            blob.upload_from_file(fobj)
+    with io.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)
 
 
 class RemoteGS(RemoteBASE):
@@ -191,21 +183,15 @@ def _upload(self, from_file, to_info, name=None, no_progress_bar=True):
     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=name or from_info.path,
-            total=blob.size,
-            bytes=True,
-            disable=no_progress_bar,
-        ) as pbar:
-            with io.open(to_file, mode="wb") as fobj:
-                raw_write = fobj.write
-
-                def write(byte_string):
-                    raw_write(byte_string)
-                    pbar.update(len(byte_string))
-
-                fobj.write = write
-                blob.download_to_file(fobj)
+        with io.open(to_file, mode="wb") as fobj:
+            with Tqdm.wrapattr(
+                fobj,
+                "write",
+                desc=name or from_info.path,
+                total=blob.size,
+                disable=no_progress_bar,
+            ) as wrapped:
+                blob.download_to_file(wrapped)
 
     def _generate_download_url(self, path_info, expires=3600):
         expiration = timedelta(seconds=int(expires))

diff --git a/setup.py b/setup.py
@@ -77,7 +77,7 @@ def run(self):
     "funcy>=1.14",
     "pathspec>=0.6.0",
     "shortuuid>=0.5.0",
-    "tqdm>=4.38.0,<5",
+    "tqdm>=4.40.0,<5",
     "packaging>=19.0",
     "win-unicode-console>=0.5; sys_platform == 'win32'",
     "pywin32>=225; sys_platform == 'win32'",