Spaces:
Runtime error
Runtime error
""" | |
Helper functionality for interoperability with stdlib `logging`. | |
""" | |
import logging | |
import sys | |
from contextlib import contextmanager | |
try: | |
from typing import Iterator, List, Optional, Type # noqa: F401 | |
except ImportError: | |
pass | |
from ..std import tqdm as std_tqdm | |
class _TqdmLoggingHandler(logging.StreamHandler): | |
def __init__( | |
self, | |
tqdm_class=std_tqdm # type: Type[std_tqdm] | |
): | |
super(_TqdmLoggingHandler, self).__init__() | |
self.tqdm_class = tqdm_class | |
def emit(self, record): | |
try: | |
msg = self.format(record) | |
self.tqdm_class.write(msg, file=self.stream) | |
self.flush() | |
except (KeyboardInterrupt, SystemExit): | |
raise | |
except: # noqa pylint: disable=bare-except | |
self.handleError(record) | |
def _is_console_logging_handler(handler): | |
return (isinstance(handler, logging.StreamHandler) | |
and handler.stream in {sys.stdout, sys.stderr}) | |
def _get_first_found_console_logging_handler(handlers): | |
for handler in handlers: | |
if _is_console_logging_handler(handler): | |
return handler | |
def logging_redirect_tqdm( | |
loggers=None, # type: Optional[List[logging.Logger]], | |
tqdm_class=std_tqdm # type: Type[std_tqdm] | |
): | |
# type: (...) -> Iterator[None] | |
""" | |
Context manager redirecting console logging to `tqdm.write()`, leaving | |
other logging handlers (e.g. log files) unaffected. | |
Parameters | |
---------- | |
loggers : list, optional | |
Which handlers to redirect (default: [logging.root]). | |
tqdm_class : optional | |
Example | |
------- | |
```python | |
import logging | |
from tqdm import trange | |
from tqdm.contrib.logging import logging_redirect_tqdm | |
LOG = logging.getLogger(__name__) | |
if __name__ == '__main__': | |
logging.basicConfig(level=logging.INFO) | |
with logging_redirect_tqdm(): | |
for i in trange(9): | |
if i == 4: | |
LOG.info("console logging redirected to `tqdm.write()`") | |
# logging restored | |
``` | |
""" | |
if loggers is None: | |
loggers = [logging.root] | |
original_handlers_list = [logger.handlers for logger in loggers] | |
try: | |
for logger in loggers: | |
tqdm_handler = _TqdmLoggingHandler(tqdm_class) | |
orig_handler = _get_first_found_console_logging_handler(logger.handlers) | |
if orig_handler is not None: | |
tqdm_handler.setFormatter(orig_handler.formatter) | |
tqdm_handler.stream = orig_handler.stream | |
logger.handlers = [ | |
handler for handler in logger.handlers | |
if not _is_console_logging_handler(handler)] + [tqdm_handler] | |
yield | |
finally: | |
for logger, original_handlers in zip(loggers, original_handlers_list): | |
logger.handlers = original_handlers | |
def tqdm_logging_redirect( | |
*args, | |
# loggers=None, # type: Optional[List[logging.Logger]] | |
# tqdm=None, # type: Optional[Type[tqdm.tqdm]] | |
**kwargs | |
): | |
# type: (...) -> Iterator[None] | |
""" | |
Convenience shortcut for: | |
```python | |
with tqdm_class(*args, **tqdm_kwargs) as pbar: | |
with logging_redirect_tqdm(loggers=loggers, tqdm_class=tqdm_class): | |
yield pbar | |
``` | |
Parameters | |
---------- | |
tqdm_class : optional, (default: tqdm.std.tqdm). | |
loggers : optional, list. | |
**tqdm_kwargs : passed to `tqdm_class`. | |
""" | |
tqdm_kwargs = kwargs.copy() | |
loggers = tqdm_kwargs.pop('loggers', None) | |
tqdm_class = tqdm_kwargs.pop('tqdm_class', std_tqdm) | |
with tqdm_class(*args, **tqdm_kwargs) as pbar: | |
with logging_redirect_tqdm(loggers=loggers, tqdm_class=tqdm_class): | |
yield pbar | |