From c7ea83d36a96e8a86b30cb056b3add5155a42ed8 Mon Sep 17 00:00:00 2001 From: Will McGugan Date: Wed, 14 Oct 2020 10:19:38 +0100 Subject: [PATCH] added progress.reset --- CHANGELOG.md | 1 + examples/listdir.py | 2 +- rich/default_styles.py | 2 +- rich/highlighter.py | 2 +- rich/progress.py | 70 +++++++++++++++++++++++++++++++-------- tests/test_highlighter.py | 9 ++++- tests/test_log.py | 2 +- tests/test_progress.py | 48 ++++++++++++++++++++++++--- 8 files changed, 113 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8f71a969..5d41fd0b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added highlighting of EUI-48 and EUI-64 (MAC addresses) - Added Console.pager - Added Console.out +- Added Progress.reset ### Changed diff --git a/examples/listdir.py b/examples/listdir.py index 52702556..bffc50af 100644 --- a/examples/listdir.py +++ b/examples/listdir.py @@ -31,5 +31,5 @@ else: ] filenames.sort(key=lambda filename: filename.lower()) filename_text = [make_filename_text(filename) for filename in filenames] - columns = Columns(filename_text, equal=True, column_first=True, expand=True) + columns = Columns(filename_text, equal=True, column_first=True) print(columns) diff --git a/rich/default_styles.py b/rich/default_styles.py index ce161437..94c106cc 100644 --- a/rich/default_styles.py +++ b/rich/default_styles.py @@ -67,7 +67,7 @@ DEFAULT_STYLES: Dict[str, Style] = { "repr.tag_name": Style(color="bright_magenta", bold=True), "repr.tag_contents": Style(color="default"), "repr.tag_end": Style(bold=True), - "repr.attrib_name": Style(color="yellow", italic=True), + "repr.attrib_name": Style(color="yellow", italic=False), "repr.attrib_equal": Style(bold=True), "repr.attrib_value": Style(color="magenta", italic=False), "repr.number": Style(color="blue", bold=True, italic=False), diff --git a/rich/highlighter.py b/rich/highlighter.py index 41ff4e91..5cfae277 100644 --- a/rich/highlighter.py +++ b/rich/highlighter.py @@ -82,13 +82,13 @@ class ReprHighlighter(RegexHighlighter): base_style = "repr." highlights = [ r"(?P\<)(?P[\w\-\.\:]*)(?P.*?)(?P\>)", + r"(?P\w+?)=(?P\"?[\w_]+\"?)?", _combine_regex( r"(?P[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})", r"(?P([A-Fa-f0-9]{1,4}::?){1,7}[A-Fa-f0-9]{1,4})", r"(?P(?:[0-9A-Fa-f]{1,2}-){7}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{1,2}:){7}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{4}\.){3}[0-9A-Fa-f]{4})", r"(?P(?:[0-9A-Fa-f]{1,2}-){5}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{1,2}:){5}[0-9A-Fa-f]{1,2}|(?:[0-9A-Fa-f]{4}\.){2}[0-9A-Fa-f]{4})", r"(?P[\{\[\(\)\]\}])", - r"(?P\w+?)=(?P\"?[\w_]+\"?)?", r"(?PTrue)|(?PFalse)|(?PNone)", r"(?P(?\B(\/[\w\.\-\_\+]+)*\/)(?P[\w\.\-\_\+]*)?", diff --git a/rich/progress.py b/rich/progress.py index b2eadc33..9541ea40 100644 --- a/rich/progress.py +++ b/rich/progress.py @@ -146,9 +146,10 @@ def track( refresh_per_second=refresh_per_second, ) - yield from progress.track( - sequence, total=total, description=description, update_period=update_period - ) + with progress: + yield from progress.track( + sequence, total=total, description=description, update_period=update_period + ) class ProgressColumn(ABC): @@ -436,6 +437,10 @@ class Task: estimate = ceil(self.remaining / speed) return estimate + def _reset(self) -> None: + """Reset progress.""" + self._progress.clear() + class _RefreshThread(Thread): """A thread that calls refresh() on the Process object at regular intervals.""" @@ -664,19 +669,19 @@ class Progress(JupyterMixin, RenderHook): task_id = self.add_task(description, total=task_total) else: self.update(task_id, total=task_total) - with self: - if self.auto_refresh: - with _TrackThread(self, task_id, update_period) as track_thread: - for value in sequence: - yield value - track_thread.completed += 1 - else: - advance = self.advance - refresh = self.refresh + + if self.auto_refresh: + with _TrackThread(self, task_id, update_period) as track_thread: for value in sequence: yield value - advance(task_id, 1) - refresh() + track_thread.completed += 1 + else: + advance = self.advance + refresh = self.refresh + for value in sequence: + yield value + advance(task_id, 1) + refresh() def start_task(self, task_id: TaskID) -> None: """Start a task. @@ -737,6 +742,7 @@ class Progress(JupyterMixin, RenderHook): if total is not None: task.total = total + task._reset() if advance is not None: task.completed += advance if completed is not None: @@ -762,6 +768,42 @@ class Progress(JupyterMixin, RenderHook): popleft() _progress.append(ProgressSample(current_time, update_completed)) + def reset( + self, + task_id: TaskID, + *, + start: bool = True, + total: Optional[int] = None, + completed: int = 0, + visible: Optional[bool] = None, + description: Optional[str] = None, + **fields: Any, + ) -> None: + """Reset a task so completed is 0 and the clock is reset. + + Args: + task_id (TaskID): ID of task. + start (bool, optional): Start the task after reset. Defaults to True. + total (int, optional): New total steps in task, or None to use current total. Defaults to None. + completed (int, optional): Number of steps completed. Defaults to 0. + **fields (str): Additional data fields required for rendering. + """ + current_time = self.get_time() + with self._lock: + task = self._tasks[task_id] + task._reset() + task.start_time = current_time if start else None + if total is not None: + task.total = total + task.completed = completed + if visible is not None: + task.visible = visible + if fields: + task.fields = fields + if description is not None: + task.description = description + self.refresh() + def advance(self, task_id: TaskID, advance: float = 1) -> None: """Advance task by a number of steps. diff --git a/tests/test_highlighter.py b/tests/test_highlighter.py index eecaa5ff..da5fc9bc 100644 --- a/tests/test_highlighter.py +++ b/tests/test_highlighter.py @@ -32,7 +32,14 @@ highlight_tests = [ ], ), ("foo=bar", [Span(0, 3, "repr.attrib_name"), Span(4, 7, "repr.attrib_value")]), - ('foo="bar"', [Span(0, 3, "repr.attrib_name"), Span(4, 9, "repr.attrib_value")]), + ( + 'foo="bar"', + [ + Span(0, 3, "repr.attrib_name"), + Span(4, 9, "repr.attrib_value"), + Span(4, 9, "repr.str"), + ], + ), ("( )", [Span(0, 1, "repr.brace"), Span(2, 3, "repr.brace")]), ("[ ]", [Span(0, 1, "repr.brace"), Span(2, 3, "repr.brace")]), ("{ }", [Span(0, 1, "repr.brace"), Span(2, 3, "repr.brace")]), diff --git a/tests/test_log.py b/tests/test_log.py index 17980cc4..5728f0fb 100644 --- a/tests/test_log.py +++ b/tests/test_log.py @@ -38,7 +38,7 @@ def render_log(): def test_log(): expected = replace_link_ids( - "\n\x1b[2;36m[TIME]\x1b[0m\x1b[2;36m \x1b[0mHello from \x1b[1m<\x1b[0m\x1b[1;95mconsole\x1b[0m\x1b[39m \x1b[0m\x1b[3;33mwidth\x1b[0m\x1b[39m=\x1b[0m\x1b[35m80\x1b[0m\x1b[39m ColorSystem.TRUECOLOR\x1b[0m\x1b[1m>\x1b[0m ! \x1b]8;id=0;foo\x1b\\\x1b[2mtest_log.py\x1b[0m\x1b]8;;\x1b\\\x1b[2m:34\x1b[0m\n\x1b[2;36m \x1b[0m\x1b[2;36m \x1b[0m\x1b[1m[\x1b[0m\x1b[1;34m1\x1b[0m, \x1b[1;34m2\x1b[0m, \x1b[1;34m3\x1b[0m\x1b[1m]\x1b[0m \x1b]8;id=0;foo\x1b\\\x1b[2mtest_log.py\x1b[0m\x1b]8;;\x1b\\\x1b[2m:35\x1b[0m\n \x1b[34m╭─\x1b[0m\x1b[34m───────────────────── \x1b[0m\x1b[3;34mlocals\x1b[0m\x1b[34m ─────────────────────\x1b[0m\x1b[34m─╮\x1b[0m \n \x1b[34m│\x1b[0m \x1b[3;33mconsole\x1b[0m\x1b[31m =\x1b[0m \x1b[1m<\x1b[0m\x1b[1;95mconsole\x1b[0m\x1b[39m \x1b[0m\x1b[3;33mwidth\x1b[0m\x1b[39m=\x1b[0m\x1b[35m80\x1b[0m\x1b[39m ColorSystem.TRUECOLOR\x1b[0m\x1b[1m>\x1b[0m \x1b[34m│\x1b[0m \n \x1b[34m╰────────────────────────────────────────────────────╯\x1b[0m \n" + "\n\x1b[2;36m[TIME]\x1b[0m\x1b[2;36m \x1b[0mHello from \x1b[1m<\x1b[0m\x1b[1;95mconsole\x1b[0m\x1b[39m \x1b[0m\x1b[33mwidth\x1b[0m\x1b[39m=\x1b[0m\x1b[1;34m80\x1b[0m\x1b[39m ColorSystem.TRUECOLOR\x1b[0m\x1b[1m>\x1b[0m ! \x1b]8;id=0;foo\x1b\\\x1b[2mtest_log.py\x1b[0m\x1b]8;;\x1b\\\x1b[2m:34\x1b[0m\n\x1b[2;36m \x1b[0m\x1b[2;36m \x1b[0m\x1b[1m[\x1b[0m\x1b[1;34m1\x1b[0m, \x1b[1;34m2\x1b[0m, \x1b[1;34m3\x1b[0m\x1b[1m]\x1b[0m \x1b]8;id=0;foo\x1b\\\x1b[2mtest_log.py\x1b[0m\x1b]8;;\x1b\\\x1b[2m:35\x1b[0m\n \x1b[34m╭─\x1b[0m\x1b[34m───────────────────── \x1b[0m\x1b[3;34mlocals\x1b[0m\x1b[34m ─────────────────────\x1b[0m\x1b[34m─╮\x1b[0m \n \x1b[34m│\x1b[0m \x1b[3;33mconsole\x1b[0m\x1b[31m =\x1b[0m \x1b[1m<\x1b[0m\x1b[1;95mconsole\x1b[0m\x1b[39m \x1b[0m\x1b[33mwidth\x1b[0m\x1b[39m=\x1b[0m\x1b[1;34m80\x1b[0m\x1b[39m ColorSystem.TRUECOLOR\x1b[0m\x1b[1m>\x1b[0m \x1b[34m│\x1b[0m \n \x1b[34m╰────────────────────────────────────────────────────╯\x1b[0m \n" ) rendered = render_log() print(repr(rendered)) diff --git a/tests/test_progress.py b/tests/test_progress.py index e94d75ed..9d4bb6ab 100644 --- a/tests/test_progress.py +++ b/tests/test_progress.py @@ -187,7 +187,15 @@ def test_track() -> None: assert value == next(expected_values) result = console.file.getvalue() print(repr(result)) - expected = "\x1b[?25ltest \x1b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m 0%\x1b[0m \x1b[36m-:--:--\x1b[0m\r\x1b[2Ktest \x1b[38;2;249;38;114m━━━━━━━━━━━━━\x1b[0m\x1b[38;5;237m╺\x1b[0m\x1b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m 33%\x1b[0m \x1b[36m-:--:--\x1b[0m\r\x1b[2Ktest \x1b[38;2;249;38;114m━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m\x1b[38;2;249;38;114m╸\x1b[0m\x1b[38;5;237m━━━━━━━━━━━━━\x1b[0m \x1b[35m 67%\x1b[0m \x1b[36m0:00:06\x1b[0m\r\x1b[2Ktest \x1b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m100%\x1b[0m \x1b[36m0:00:00\x1b[0m\r\x1b[2Ktest \x1b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m100%\x1b[0m \x1b[36m0:00:00\x1b[0m\n\x1b[?25h" + expected = "\x1b[?25l\r\x1b[2Ktest \x1b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m 0%\x1b[0m \x1b[36m-:--:--\x1b[0m\r\x1b[2Ktest \x1b[38;2;249;38;114m━━━━━━━━━━━━━\x1b[0m\x1b[38;5;237m╺\x1b[0m\x1b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m 33%\x1b[0m \x1b[36m-:--:--\x1b[0m\r\x1b[2Ktest \x1b[38;2;249;38;114m━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m\x1b[38;2;249;38;114m╸\x1b[0m\x1b[38;5;237m━━━━━━━━━━━━━\x1b[0m \x1b[35m 67%\x1b[0m \x1b[36m0:00:06\x1b[0m\r\x1b[2Ktest \x1b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m100%\x1b[0m \x1b[36m0:00:00\x1b[0m\r\x1b[2Ktest \x1b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m100%\x1b[0m \x1b[36m0:00:00\x1b[0m\n\x1b[?25h" + print("--") + print("RESULT:") + print(result) + print(repr(result)) + print("EXPECTED:") + print(expected) + print(repr(expected)) + assert result == expected with pytest.raises(ValueError): @@ -208,11 +216,18 @@ def test_progress_track() -> None: ) test = ["foo", "bar", "baz"] expected_values = iter(test) - for value in progress.track(test, description="test"): - assert value == next(expected_values) + with progress: + for value in progress.track(test, description="test"): + assert value == next(expected_values) result = console.file.getvalue() print(repr(result)) - expected = "\x1b[?25ltest \x1b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m 0%\x1b[0m \x1b[36m-:--:--\x1b[0m\r\x1b[2Ktest \x1b[38;2;249;38;114m━━━━━━━━━━━━━\x1b[0m\x1b[38;5;237m╺\x1b[0m\x1b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m 33%\x1b[0m \x1b[36m-:--:--\x1b[0m\r\x1b[2Ktest \x1b[38;2;249;38;114m━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m\x1b[38;2;249;38;114m╸\x1b[0m\x1b[38;5;237m━━━━━━━━━━━━━\x1b[0m \x1b[35m 67%\x1b[0m \x1b[36m0:00:06\x1b[0m\r\x1b[2Ktest \x1b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m100%\x1b[0m \x1b[36m0:00:00\x1b[0m\r\x1b[2Ktest \x1b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m100%\x1b[0m \x1b[36m0:00:00\x1b[0m\n\x1b[?25h" + expected = "\x1b[?25l\r\x1b[2Ktest \x1b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m 0%\x1b[0m \x1b[36m-:--:--\x1b[0m\r\x1b[2Ktest \x1b[38;2;249;38;114m━━━━━━━━━━━━━\x1b[0m\x1b[38;5;237m╺\x1b[0m\x1b[38;5;237m━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m 33%\x1b[0m \x1b[36m-:--:--\x1b[0m\r\x1b[2Ktest \x1b[38;2;249;38;114m━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m\x1b[38;2;249;38;114m╸\x1b[0m\x1b[38;5;237m━━━━━━━━━━━━━\x1b[0m \x1b[35m 67%\x1b[0m \x1b[36m0:00:06\x1b[0m\r\x1b[2Ktest \x1b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m100%\x1b[0m \x1b[36m0:00:00\x1b[0m\r\x1b[2Ktest \x1b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\x1b[0m \x1b[35m100%\x1b[0m \x1b[36m0:00:00\x1b[0m\n\x1b[?25h" + + print(expected) + print(repr(expected)) + print(result) + print(repr(result)) + assert result == expected with pytest.raises(ValueError): @@ -323,6 +338,31 @@ def test_track_thread() -> None: track_thread.completed += 1 +def test_reset() -> None: + progress = Progress() + task_id = progress.add_task("foo") + progress.advance(task_id, 1) + progress.advance(task_id, 1) + progress.advance(task_id, 1) + progress.advance(task_id, 7) + task = progress.tasks[task_id] + assert task.completed == 10 + progress.reset( + task_id, + total=200, + completed=20, + visible=False, + description="bar", + example="egg", + ) + assert task.total == 200 + assert task.completed == 20 + assert task.visible == False + assert task.description == "bar" + assert task.fields == {"example": "egg"} + assert not task._progress + + if __name__ == "__main__": _render = render_progress() print(_render)