Rich 控制台输出美化库

学习

rich/live.py 是实现异步自动刷新渲染的关键模块。

1. 自动刷新线程：

# 继承线程类，覆写run方法为刷新逻辑，使用 Event 同步
class _RefreshThread(Thread):
    """A thread that calls refresh() at regular intervals."""

    def __init__(self, live: "Live", refresh_per_second: float) -> None:
        self.live = live
        self.refresh_per_second = refresh_per_second
        self.done = Event()
        super().__init__(daemon=True)

		# 当其他线程调用该线程的该方法时，在 run
    def stop(self) -> None:
        self.done.set()

    def run(self) -> None:
		    # 每一帧都去询问 event 是否 set，是则停止刷新动画
        while not self.done.wait(1 / self.refresh_per_second):
            with self.live._lock:
                if not self.done.is_set():
                    self.live.refresh()

• 使用 threading 的 Thread 类实现多线程，Event 类实现基于事件的线程同步。
• Event 类相当于简单的带超时的 Condition 锁（基于 Producer-Consumer 同步模型）。

1. Live 的实现：

class Live():
...
     def start(self, refresh: bool = False) -> None:
        ...
          if self.auto_refresh:
              self._refresh_thread = _RefreshThread(self, self.refresh_per_second)
              self._refresh_thread.start()
              
     def stop(self) -> None:
        with self._lock:
            if self.auto_refresh and self._refresh_thread is not None:
                self._refresh_thread.stop()
                self._refresh_thread = None

1. 上下文管理包装 Live 实现：

class Status(JupyterMixin):
    def start(self) -> None:
        """Start the status animation."""
        self._live.start()

    def stop(self) -> None:
        """Stop the spinner animation."""
        self._live.stop()

    def __enter__(self) -> "Status":
        self.start()
        return self

    def __exit__(self, ) -> None:
        self.stop()
 
 # 使用方法       
 with Status():
		    time.sleep(60) # 退出上下文前RefreshThread会一直刷新动画

思考

• Status 是怎么保证刷新动画一直位于控制台底层的？
  • 更多控制台转义字符详见https://zh.wikipedia.org/wiki/ANSI%E8%BD%AC%E4%B9%89%E5%BA%8F%E5%88%97

实现

import sys
import time
from threading import Thread, Lock, Event

class Refreshable():
    def refresh(self):
        pass

class RefreshThread(Thread):
    """Refresh Thread"""
    def __init__(self, live: Refreshable, rate=3):
        """
            live (Refreshable): Thread-safe Refresh Resource.
            rate (int, optional): Refresh Rate. Defaults to 2.
        """
        self._live = live
        self._done = Event()
        self._rate = rate
        super().__init__(daemon=True)

    def stop(self):
        self._done.set()

    def run(self):
        while not self._done.wait(1 / self._rate):
            self._live.refresh()

class VizIter(Refreshable):
    """
    Thread-safe Visualize iteration progress. Each time this object being iterated, a message will be print to screen for progress information.
    Useful in polling loop.
    """

    def __init__(
        self,
        desc: str = "",
        max_cnt: int = sys.maxsize,
        iter_chars: tuple[str] = ('\\\\', '|', '/', '-'),
    ):
        """Create a new Bar

        Args:
            desc (str, optional): Description message. Defaults to None.
            iter_chars (list[str], optional): Iterable ascii chars.
            max_cnt (int, optional): max counts. Defaults to infinite.
        """
        self.desc = desc
        self._chars = iter_chars
        self._len = len(self._chars)
        self._max = max_cnt
        self._cnt = 0
        self._lock = Lock()
        self._thread = None

    def __next__(self):
        """counter +1"""
        self._lock.acquire()
        if self._cnt == self._max:
            raise StopIteration
        print(
            (f"{self.desc} {self._chars[self._cnt % self._len]}"),
            end="\\r",
        )
        ret = self._cnt
        self._cnt = self._cnt + 1
        self._lock.release()
        return ret

    def __iter__(self):
        return self

    def __enter__(self):
        self._thread = RefreshThread(self)
        self._thread.start()
        return self

    def __exit__(self, _type, _value, traceback):
        self._thread.stop()
        self._thread = None

    def refresh(self):
        next(self)

    def msg(self, desc: str):
        """Change message shown"""
        self._lock.acquire()
        self.desc = desc
        self._lock.release()
        return self

with VizIter("Waiting...") as progress:
    num = 5
    for idx, task in enumerate([f"Task{i} has been compeleted!" for i in range(num)]):
        time.sleep(3)
        print(task)
        progress.msg("Waiting... [%d/%d]" % (idx, num))