127 lines
3.7 KiB
Python
127 lines
3.7 KiB
Python
"""
|
|
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().__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
|
|
|
|
|
|
@contextmanager
|
|
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
|
|
|
|
|
|
@contextmanager
|
|
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
|