28 lines
762 B
Python
28 lines
762 B
Python
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])
|