Add Rich Progressbar Wrapper

This commit is contained in:
Matthias 2023-04-09 16:04:31 +02:00
parent df51111c33
commit d532da9071

View File

@ -0,0 +1,27 @@
import logging
import sys
from contextlib import contextmanager
from rich.progress import Progress
@contextmanager
def FtProgress(*args, **kwargs):
"""
Wrapper around rich.progress.Progress to fix issues with logging.
"""
try:
__logger = kwargs.pop('logger', None)
streamhandlers = [x for x in __logger.root.handlers if type(x) == logging.StreamHandler]
__prior_stderr = []
with Progress(*args, **kwargs) as progress:
for handler in streamhandlers:
__prior_stderr.append(handler.stream)
handler.setStream(sys.stderr)
yield progress
finally:
for idx, handler in enumerate(streamhandlers):
handler.setStream(__prior_stderr[idx])