本篇文章介绍airbus中进度条装饰器,对于有循环次数的函数,可以使用该装饰器来比整个任务添加一个进度条监控。
1. 源码
from functools import wrapsfrom inspect import isgeneratorfrom typing import Textfrom rich.panel import Panelfrom rich.progress import Progress, BarColumn, SpinnerColumn, TimeRemainingColumn, TimeElapsedColumnclass 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.valuereturn wrapper
2. 使用方法
如果需要使用进度条装饰器的话,需要所修饰的函数为生成器函数,并且yield返回第一个值为循环次数:
import timefrom airbus.decorators import progressbar@progressbardef dummy_loop():total = 20yield totalfor i in range(1, total + 1):yield itime.sleep(1)return "done"print(dummy_loop())
