• 1. 源码
• 2. 使用方法

本篇文章介绍airbus中进度条装饰器，对于有循环次数的函数，可以使用该装饰器来比整个任务添加一个进度条监控。

1. 源码

from functools import wraps
from inspect import isgenerator
from typing import Text
from rich.panel import Panel
from rich.progress import Progress, BarColumn, SpinnerColumn, TimeRemainingColumn, TimeElapsedColumn
class RichProgress(Progress):
    def get_renderables(self):
        yield Panel(self.make_tasks_table(self.tasks), title="⌛ Progress Frame", expand=False)
def progressbar(func):
    """Add a progress bar decorator for generator function to display in the terminal"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        func_generator = func(*args, **kwargs)
        if not isgenerator(func_generator):
            raise Exception(f"{func.__name__} is not a generator function!")
        total = next(func_generator)
        with RichProgress(
                "[progress.description]{task.description}({task.completed}/{task.total})",
                SpinnerColumn(finished_text="🚀"),
                BarColumn(),
                "[progress.percentage]{task.percentage:>3.2f}%",
                SpinnerColumn(spinner_name="clock", finished_text="🕐"),
                TimeElapsedColumn(),
                "⏳",
                TimeRemainingColumn()) as progress:
            description = "[red]Loading"
            task = progress.add_task(description, total=total)
            try:
                while True:
                    p = next(func_generator)
                    if p != total:
                        description = "[red]Loading"
                    else:
                        description = "[green]Finished"
                    progress.update(task, completed=p, description=description)
            except StopIteration as result:
                return result.value
    return wrapper

2. 使用方法

如果需要使用进度条装饰器的话，需要所修饰的函数为生成器函数，并且yield返回第一个值为循环次数：

import time
from airbus.decorators import progressbar
@progressbar
def dummy_loop():
    total = 20
    yield total
    for i in range(1, total + 1):
        yield i
        time.sleep(1)
    return "done"
print(dummy_loop())