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

1. 源码

  1. from functools import wraps
  2. from inspect import isgenerator
  3. from typing import Text
  4. from rich.panel import Panel
  5. from rich.progress import Progress, BarColumn, SpinnerColumn, TimeRemainingColumn, TimeElapsedColumn
  6. class RichProgress(Progress):
  7. def get_renderables(self):
  8. yield Panel(self.make_tasks_table(self.tasks), title="⌛ Progress Frame", expand=False)
  9. def progressbar(func):
  10. """Add a progress bar decorator for generator function to display in the terminal"""
  11. @wraps(func)
  12. def wrapper(*args, **kwargs):
  13. func_generator = func(*args, **kwargs)
  14. if not isgenerator(func_generator):
  15. raise Exception(f"{func.__name__} is not a generator function!")
  16. total = next(func_generator)
  17. with RichProgress(
  18. "[progress.description]{task.description}({task.completed}/{task.total})",
  19. SpinnerColumn(finished_text="🚀"),
  20. BarColumn(),
  21. "[progress.percentage]{task.percentage:>3.2f}%",
  22. SpinnerColumn(spinner_name="clock", finished_text="🕐"),
  23. TimeElapsedColumn(),
  24. "⏳",
  25. TimeRemainingColumn()) as progress:
  26. description = "[red]Loading"
  27. task = progress.add_task(description, total=total)
  28. try:
  29. while True:
  30. p = next(func_generator)
  31. if p != total:
  32. description = "[red]Loading"
  33. else:
  34. description = "[green]Finished"
  35. progress.update(task, completed=p, description=description)
  36. except StopIteration as result:
  37. return result.value
  38. return wrapper

2. 使用方法

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

  1. import time
  2. from airbus.decorators import progressbar
  3. @progressbar
  4. def dummy_loop():
  5. total = 20
  6. yield total
  7. for i in range(1, total + 1):
  8. yield i
  9. time.sleep(1)
  10. return "done"
  11. print(dummy_loop())