From 12961b35d609e15f51526429afe997315596a8a7 Mon Sep 17 00:00:00 2001
From: Roman Donchenko <roman@cvat.ai>
Date: Mon, 7 Nov 2022 07:02:09 +0300
Subject: [PATCH] cvat-sdk: make sure to redraw the progress bar after use
 (#5235)

tqdm does not necessarily redraw the progress bar after every update,
and without the call to `refresh` it might not update it to the final state.

I discovered this while writing a test in which the progress bar is updated
in chunks, and the last chunk is smaller than the previous ones.
---
 cvat-sdk/cvat_sdk/core/helpers.py   |  9 +++++++++
 cvat-sdk/cvat_sdk/core/uploading.py | 13 +++++++------
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/cvat-sdk/cvat_sdk/core/helpers.py b/cvat-sdk/cvat_sdk/core/helpers.py
index 6dda3b681..b548d85cd 100644
--- a/cvat-sdk/cvat_sdk/core/helpers.py
+++ b/cvat-sdk/cvat_sdk/core/helpers.py
@@ -60,6 +60,9 @@ class TqdmProgressReporter(ProgressReporter):
     def advance(self, delta: int):
         self.tqdm.update(delta)
 
+    def finish(self):
+        self.tqdm.refresh()
+
 
 class StreamWithProgress:
     def __init__(self, stream: io.RawIOBase, pbar: ProgressReporter, length: Optional[int] = None):
@@ -88,6 +91,12 @@ class StreamWithProgress:
     def tell(self):
         return self.stream.tell()
 
+    def __enter__(self) -> StreamWithProgress:
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback) -> None:
+        self.pbar.finish()
+
 
 def expect_status(codes: Union[int, Iterable[int]], response: urllib3.HTTPResponse) -> None:
     if not hasattr(codes, "__iter__"):
diff --git a/cvat-sdk/cvat_sdk/core/uploading.py b/cvat-sdk/cvat_sdk/core/uploading.py
index 3bfe630ff..d51fa76d7 100644
--- a/cvat-sdk/cvat_sdk/core/uploading.py
+++ b/cvat-sdk/cvat_sdk/core/uploading.py
@@ -15,7 +15,7 @@ import urllib3
 from cvat_sdk.api_client.api_client import ApiClient, Endpoint
 from cvat_sdk.api_client.rest import RESTClientObject
 from cvat_sdk.core.helpers import StreamWithProgress, expect_status
-from cvat_sdk.core.progress import ProgressReporter
+from cvat_sdk.core.progress import NullProgressReporter, ProgressReporter
 
 if TYPE_CHECKING:
     from cvat_sdk.core.client import Client
@@ -237,16 +237,17 @@ class Uploader:
         CHUNK_SIZE = 10 * 2**20
 
         file_size = os.stat(filename).st_size
+        if pbar is None:
+            pbar = NullProgressReporter()
 
-        with open(filename, "rb") as input_file:
-            if pbar is not None:
-                input_file = StreamWithProgress(input_file, pbar, length=file_size)
-
+        with open(filename, "rb") as input_file, StreamWithProgress(
+            input_file, pbar, length=file_size
+        ) as input_file_with_progress:
             tus_uploader = self._make_tus_uploader(
                 self._client.api_client,
                 url=url.rstrip("/") + "/",
                 metadata=meta,
-                file_stream=input_file,
+                file_stream=input_file_with_progress,
                 chunk_size=CHUNK_SIZE,
                 log_func=logger,
             )
-- 
GitLab