codekingpro/portable-devtools

1"""2Helper functionality for interoperability with stdlib `logging`.3"""4import logging5import sys6from contextlib import contextmanager7 8try:9    from typing import Iterator, List, Optional, Type  # noqa: F40110except ImportError:11    pass12 13from ..std import tqdm as std_tqdm14 15 16class _TqdmLoggingHandler(logging.StreamHandler):17    def __init__(18        self,19        tqdm_class=std_tqdm  # type: Type[std_tqdm]20    ):21        super().__init__()22        self.tqdm_class = tqdm_class23 24    def emit(self, record):25        try:26            msg = self.format(record)27            self.tqdm_class.write(msg, file=self.stream)28            self.flush()29        except (KeyboardInterrupt, SystemExit):30            raise31        except:  # noqa pylint: disable=bare-except32            self.handleError(record)33 34 35def _is_console_logging_handler(handler):36    return (isinstance(handler, logging.StreamHandler)37            and handler.stream in {sys.stdout, sys.stderr})38 39 40def _get_first_found_console_logging_handler(handlers):41    for handler in handlers:42        if _is_console_logging_handler(handler):43            return handler44 45 46@contextmanager47def logging_redirect_tqdm(48    loggers=None,  # type: Optional[List[logging.Logger]],49    tqdm_class=std_tqdm  # type: Type[std_tqdm]50):51    # type: (...) -> Iterator[None]52    """53    Context manager redirecting console logging to `tqdm.write()`, leaving54    other logging handlers (e.g. log files) unaffected.55 56    Parameters57    ----------58    loggers  : list, optional59      Which handlers to redirect (default: [logging.root]).60    tqdm_class  : optional61 62    Example63    -------64    ```python65    import logging66    from tqdm import trange67    from tqdm.contrib.logging import logging_redirect_tqdm68 69    LOG = logging.getLogger(__name__)70 71    if __name__ == '__main__':72        logging.basicConfig(level=logging.INFO)73        with logging_redirect_tqdm():74            for i in trange(9):75                if i == 4:76                    LOG.info("console logging redirected to `tqdm.write()`")77        # logging restored78    ```79    """80    if loggers is None:81        loggers = [logging.root]82    original_handlers_list = [logger.handlers for logger in loggers]83    try:84        for logger in loggers:85            tqdm_handler = _TqdmLoggingHandler(tqdm_class)86            orig_handler = _get_first_found_console_logging_handler(logger.handlers)87            if orig_handler is not None:88                tqdm_handler.setFormatter(orig_handler.formatter)89                tqdm_handler.stream = orig_handler.stream90            logger.handlers = [91                handler for handler in logger.handlers92                if not _is_console_logging_handler(handler)] + [tqdm_handler]93        yield94    finally:95        for logger, original_handlers in zip(loggers, original_handlers_list):96            logger.handlers = original_handlers97 98 99@contextmanager100def tqdm_logging_redirect(101    *args,102    # loggers=None,  # type: Optional[List[logging.Logger]]103    # tqdm=None,  # type: Optional[Type[tqdm.tqdm]]104    **kwargs105):106    # type: (...) -> Iterator[None]107    """108    Convenience shortcut for:109    ```python110    with tqdm_class(*args, **tqdm_kwargs) as pbar:111        with logging_redirect_tqdm(loggers=loggers, tqdm_class=tqdm_class):112            yield pbar113    ```114 115    Parameters116    ----------117    tqdm_class  : optional, (default: tqdm.std.tqdm).118    loggers  : optional, list.119    **tqdm_kwargs  : passed to `tqdm_class`.120    """121    tqdm_kwargs = kwargs.copy()122    loggers = tqdm_kwargs.pop('loggers', None)123    tqdm_class = tqdm_kwargs.pop('tqdm_class', std_tqdm)124    with tqdm_class(*args, **tqdm_kwargs) as pbar:125        with logging_redirect_tqdm(loggers=loggers, tqdm_class=tqdm_class):126            yield pbar127