本篇文章介绍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())