diff --git a/gnes/client/cli.py b/gnes/client/cli.py index b44d4d22..ada594ec 100644 --- a/gnes/client/cli.py +++ b/gnes/client/cli.py @@ -17,7 +17,7 @@ import sys import time import zipfile -from typing import List, Generator +from typing import Generator from termcolor import colored @@ -62,22 +62,8 @@ def query_callback(self, req: 'gnes_pb2.Request', resp: 'gnes_pb2.Response'): print(req) print(resp) - def read_all(self) -> List[bytes]: - if self.args.txt_file: - all_bytes = [v.encode() for v in self.args.txt_file] - elif self.args.image_zip_file: - zipfile_ = zipfile.ZipFile(self.args.image_zip_file) - all_bytes = [zipfile_.open(v).read() for v in zipfile_.namelist()] - elif self.args.video_zip_file: - zipfile_ = zipfile.ZipFile(self.args.video_zip_file) - all_bytes = [zipfile_.open(v).read() for v in zipfile_.namelist()] - else: - raise AttributeError('--txt_file, --image_zip_file, --video_zip_file one must be given') - - return all_bytes - @property - def bytes_generator(self) -> Generator[bytes]: + def bytes_generator(self) -> Generator[bytes, None, None]: if self.args.txt_file: all_bytes = (v.encode() for v in self.args.txt_file) elif self.args.image_zip_file: @@ -107,7 +93,7 @@ def update(self): self.num_bars -= self.bar_len sys.stdout.write('\n') sys.stdout.write( - '{:>10} [{:<{}}] {:3.0f}% {:>8}: {:3.1f}s {:>8}: {:3.1f} batch/s'.format( + '{:>10} [{:<{}}] {:>8}: {:3.1f}s {:>8}: {:3.1f} batch/s'.format( colored(self.task_name, 'cyan'), colored('=' * self.num_bars, 'green'), self.bar_len + 9, diff --git a/tests/test_progressbar.py b/tests/test_progressbar.py new file mode 100644 index 00000000..53177f98 --- /dev/null +++ b/tests/test_progressbar.py @@ -0,0 +1,22 @@ +import time +import unittest + +from gnes.client.cli import ProgressBar + + +class TestProgessbar(unittest.TestCase): + def setUp(self): + self.bar_len = 20 + + def test_progressbar5(self): + # should be 5 line + with ProgressBar(task_name='test', bar_len=self.bar_len) as pb: + for j in range(5 * self.bar_len): + pb.update() + + def test_progressbar1(self): + # should be single line + with ProgressBar(task_name='test', bar_len=self.bar_len) as pb: + for j in range(self.bar_len): + pb.update() + time.sleep(.1)