未验证 提交 12961b35 编写于 作者: R Roman Donchenko 提交者: GitHub

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.
上级 16406fec
......@@ -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__"):
......
......@@ -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,
)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册