diff --git a/README.rst b/README.rst index dd7e5a381..0ea0e98b0 100644 --- a/README.rst +++ b/README.rst @@ -1321,6 +1321,59 @@ A reusable canonical example is given below: # After the `with`, printing is restored print("Done!") +Redirecting console logging to tqdm +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Similar to redirecting ``sys.stdout`` directly as detailed in the previous section, +you may want to redirect logging that would otherwise go to the +console (``sys.stdout`` or ``sys.stderr``) to ``tqdm``. + +Note: if you are also replace ``sys.stdout`` and ``sys.stderr`` at the same time, +then the logging should be redirected first. Otherwise it won't be able to detect +the console logging handler. + +For that you may use ``redirect_logging_to_tqdm`` or ``tqdm_with_logging_redirect`` +from ``tqdm.contrib.logging``. Both methods accept the following optional parameters: + +- ``loggers``: A list of loggers to update. Defaults to ``logging.root``. +- ``tqdm``: A ``tqdm`` class. Defaults to ``tqdm.tqdm``. + +An example redirecting the console logging to tqdm: + +.. code:: python + + import logging + from tqdm.contrib.logging import redirect_logging_to_tqdm + + LOGGER = logging.getLogger(__name__) + + if __name__ == '__main__': + logging.basicConfig(level='INFO') + with redirect_logging_to_tqdm(): + # logging to the console is now redirected to tqdm + LOGGER.info('some message') + # logging is now restored + +An similar example, wrapping tqdm while redirecting console logging: + +.. code:: python + + import logging + from tqdm.contrib.logging import tqdm_with_logging_redirect + + LOGGER = logging.getLogger(__name__) + + if __name__ == '__main__': + logging.basicConfig(level='INFO') + + file_list = ['file1', 'file2'] + with tqdm_with_logging_redirect(total=len(file_list)) as pbar: + # logging to the console is now redirected to tqdm + for filename in file_list: + LOGGER.info('processing file: %s', filename) + pbar.update(1) + # logging is now restored + Monitoring thread, intervals and miniters ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/tests/contrib/__init__.py b/tests/contrib/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/contrib/tests_logging.py b/tests/contrib/tests_logging.py new file mode 100644 index 000000000..0b3342257 --- /dev/null +++ b/tests/contrib/tests_logging.py @@ -0,0 +1,182 @@ +# pylint: disable=missing-module-docstring, missing-class-docstring +# pylint: disable=missing-function-docstring, no-self-use + +from __future__ import absolute_import + +import logging +import logging.handlers +import sys +from io import StringIO + +import pytest + +from tqdm import tqdm +from tqdm.contrib.logging import _get_first_found_console_logging_formatter +from tqdm.contrib.logging import _TqdmLoggingHandler as TqdmLoggingHandler +from tqdm.contrib.logging import redirect_logging_to_tqdm, tqdm_with_logging_redirect + +from ..tests_tqdm import importorskip + +LOGGER = logging.getLogger(__name__) + +TEST_LOGGING_FORMATTER = logging.Formatter() + + +class CustomTqdm(tqdm): + messages = [] + + @classmethod + def write(cls, s, **__): # pylint: disable=arguments-differ + CustomTqdm.messages.append(s) + + +class ErrorRaisingTqdm(tqdm): + exception_class = RuntimeError + + @classmethod + def write(cls, s, **__): # pylint: disable=arguments-differ + raise ErrorRaisingTqdm.exception_class('fail fast') + + +class TestTqdmLoggingHandler: + def test_should_call_tqdm_write(self): + CustomTqdm.messages = [] + logger = logging.Logger('test') + logger.handlers = [TqdmLoggingHandler(CustomTqdm)] + logger.info('test') + assert CustomTqdm.messages == ['test'] + + def test_should_call_handle_error_if_exception_was_thrown(self): + patch = importorskip('unittest.mock').patch + logger = logging.Logger('test') + ErrorRaisingTqdm.exception_class = RuntimeError + handler = TqdmLoggingHandler(ErrorRaisingTqdm) + logger.handlers = [handler] + with patch.object(handler, 'handleError') as mock: + logger.info('test') + assert mock.called + + @pytest.mark.parametrize('exception_class', [ + KeyboardInterrupt, + SystemExit + ]) + def test_should_not_swallow_certain_exceptions(self, exception_class): + logger = logging.Logger('test') + ErrorRaisingTqdm.exception_class = exception_class + handler = TqdmLoggingHandler(ErrorRaisingTqdm) + logger.handlers = [handler] + with pytest.raises(exception_class): + logger.info('test') + + +class TestGetFirstFoundConsoleLoggingFormatter: + def test_should_return_none_for_no_handlers(self): + assert _get_first_found_console_logging_formatter([]) is None + + def test_should_return_none_without_stream_handler(self): + handler = logging.handlers.MemoryHandler(capacity=1) + handler.formatter = TEST_LOGGING_FORMATTER + assert _get_first_found_console_logging_formatter([handler]) is None + + def test_should_return_none_for_stream_handler_not_stdout_or_stderr(self): + handler = logging.StreamHandler(StringIO()) + handler.formatter = TEST_LOGGING_FORMATTER + assert _get_first_found_console_logging_formatter([handler]) is None + + def test_should_return_stream_handler_formatter_if_stream_is_stdout(self): + handler = logging.StreamHandler(sys.stdout) + handler.formatter = TEST_LOGGING_FORMATTER + assert _get_first_found_console_logging_formatter( + [handler] + ) == TEST_LOGGING_FORMATTER + + def test_should_return_stream_handler_formatter_if_stream_is_stderr(self): + handler = logging.StreamHandler(sys.stderr) + handler.formatter = TEST_LOGGING_FORMATTER + assert _get_first_found_console_logging_formatter( + [handler] + ) == TEST_LOGGING_FORMATTER + + +class TestRedirectLoggingToTqdm: + def test_should_add_and_remove_tqdm_handler(self): + logger = logging.Logger('test') + with redirect_logging_to_tqdm(loggers=[logger]): + assert len(logger.handlers) == 1 + assert isinstance(logger.handlers[0], TqdmLoggingHandler) + assert not logger.handlers + + def test_should_remove_and_restore_console_handlers(self): + logger = logging.Logger('test') + stderr_console_handler = logging.StreamHandler(sys.stderr) + stdout_console_handler = logging.StreamHandler(sys.stderr) + logger.handlers = [stderr_console_handler, stdout_console_handler] + with redirect_logging_to_tqdm(loggers=[logger]): + assert len(logger.handlers) == 1 + assert isinstance(logger.handlers[0], TqdmLoggingHandler) + assert logger.handlers == [stderr_console_handler, stdout_console_handler] + + def test_should_inherit_console_logger_formatter(self): + logger = logging.Logger('test') + formatter = logging.Formatter('custom: %(message)s') + console_handler = logging.StreamHandler(sys.stderr) + console_handler.setFormatter(formatter) + logger.handlers = [console_handler] + with redirect_logging_to_tqdm(loggers=[logger]): + assert logger.handlers[0].formatter == formatter + + def test_should_not_remove_stream_handlers_not_fot_stdout_or_stderr(self): + logger = logging.Logger('test') + stream_handler = logging.StreamHandler(StringIO()) + logger.addHandler(stream_handler) + with redirect_logging_to_tqdm(loggers=[logger]): + assert len(logger.handlers) == 2 + assert logger.handlers[0] == stream_handler + assert isinstance(logger.handlers[1], TqdmLoggingHandler) + assert logger.handlers == [stream_handler] + + +class TestTqdmWithLoggingRedirect: + def test_should_add_and_remove_handler_from_root_logger_by_default(self): + original_handlers = list(logging.root.handlers) + with tqdm_with_logging_redirect(total=1) as pbar: + assert isinstance(logging.root.handlers[-1], TqdmLoggingHandler) + LOGGER.info('test') + pbar.update(1) + assert logging.root.handlers == original_handlers + + def test_should_add_and_remove_handler_from_custom_logger(self): + logger = logging.Logger('test') + with tqdm_with_logging_redirect(total=1, loggers=[logger]) as pbar: + assert len(logger.handlers) == 1 + assert isinstance(logger.handlers[0], TqdmLoggingHandler) + logger.info('test') + pbar.update(1) + assert not logger.handlers + + def test_should_not_fail_with_logger_without_console_handler(self): + logger = logging.Logger('test') + logger.handlers = [] + with tqdm_with_logging_redirect(total=1, loggers=[logger]): + logger.info('test') + assert not logger.handlers + + def test_should_format_message(self): + logger = logging.Logger('test') + console_handler = logging.StreamHandler(sys.stdout) + console_handler.setFormatter(logging.Formatter( + r'prefix:%(message)s' + )) + logger.handlers = [console_handler] + CustomTqdm.messages = [] + with tqdm_with_logging_redirect(loggers=[logger], tqdm=CustomTqdm): + logger.info('test') + assert CustomTqdm.messages == ['prefix:test'] + + def test_use_root_logger_by_default_and_write_to_custom_tqdm(self): + logger = logging.root + CustomTqdm.messages = [] + with tqdm_with_logging_redirect(total=1, tqdm=CustomTqdm) as pbar: + assert isinstance(pbar, CustomTqdm) + logger.info('test') + assert CustomTqdm.messages == ['test'] diff --git a/tqdm/contrib/logging.py b/tqdm/contrib/logging.py new file mode 100644 index 000000000..08c9c3c40 --- /dev/null +++ b/tqdm/contrib/logging.py @@ -0,0 +1,194 @@ + +""" +Enables multiple commonly used features relating to logging +in combination with tqdm. +""" +from __future__ import absolute_import + +import logging +import sys +from contextlib import contextmanager + +try: + from typing import Iterator, List, Optional, Type # pylint: disable=unused-import +except ImportError: + # we may ignore type hints + pass + +from ..std import tqdm as _tqdm + + +class _TqdmLoggingHandler(logging.StreamHandler): + def __init__( + self, + tqdm=None # type: Optional[Type[tqdm.tqdm]] + ): + super( # pylint: disable=super-with-arguments + _TqdmLoggingHandler, self + ).__init__() + if tqdm is None: + tqdm = _tqdm + self.tqdm = tqdm + + def emit(self, record): + try: + msg = self.format(record) + self.tqdm.write(msg) + 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_formatter(handlers): + for handler in handlers: + if _is_console_logging_handler(handler): + return handler.formatter + return None + + +@contextmanager +def redirect_logging_to_tqdm( + loggers=None, # type: Optional[List[logging.Logger]], + tqdm=None # type: Optional[Type[tqdm.tqdm]] +): + # type: (...) -> Iterator[None] + """ + Context manager for redirecting logging console output to tqdm. + Logging to other logging handlers, such as a log file, + will not be affected. + + By default the, the handlers of the root logger will be amended. + (for the duration of the context) + You may also provide a list of `loggers` instead + (e.g. if a particular logger doesn't fallback to the root logger) + + Example: + + ```python + import logging + from tqdm.contrib.logging import redirect_logging_to_tqdm + + LOGGER = logging.getLogger(__name__) + + if __name__ == '__main__': + logging.basicConfig(level='INFO') + with redirect_logging_to_tqdm(): + # logging to the console is now redirected to tqdm + LOGGER.info('some message') + # logging is now 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) + tqdm_handler.setFormatter( + _get_first_found_console_logging_formatter( + logger.handlers + ) + ) + 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 _pop_optional( + kwargs, # type: dict + key, # type: str + default_value=None +): + try: + return kwargs.pop(key) + except KeyError: + return default_value + + +@contextmanager +def tqdm_with_logging_redirect( + *args, + # loggers=None, # type: Optional[List[logging.Logger]] + # tqdm=None, # type: Optional[Type[tqdm.tqdm]] + **kwargs +): + # type: (...) -> Iterator[None] + """ + Similar to `redirect_logging_to_tqdm`, + but provides a context manager wrapping tqdm. + + All parameters, except `loggers` and `tqdm`, will get passed on to `tqdm`. + + By default this will wrap `tqdm.tqdm`. + You may pass your own `tqdm` class if desired. + + Example: + + ```python + import logging + from tqdm.contrib.logging import tqdm_with_logging_redirect + + LOGGER = logging.getLogger(__name__) + + if __name__ == '__main__': + logging.basicConfig(level='INFO') + + file_list = ['file1', 'file2'] + with tqdm_with_logging_redirect(total=len(file_list)) as pbar: + # logging to the console is now redirected to tqdm + for filename in file_list: + LOGGER.info('processing file: %s', filename) + pbar.update(1) + # logging is now restored + ``` + + A more advanced example with non-default tqdm class and loggers: + + ```python + import logging + from tqdm.auto import tqdm + from tqdm.contrib.logging import tqdm_with_logging_redirect + + LOGGER = logging.getLogger(__name__) + + if __name__ == '__main__': + logging.basicConfig(level='INFO') + + file_list = ['file1', 'file2'] + with tqdm_with_logging_redirect( + total=len(file_list), + tqdm=tqdm, + loggers=[LOGGER] + ) as pbar: + # logging to the console is now redirected to tqdm + for filename in file_list: + LOGGER.info('processing file: %s', filename) + pbar.update(1) + # logging is now restored + ``` + + """ + loggers = _pop_optional(kwargs, 'loggers') + tqdm = _pop_optional(kwargs, 'tqdm') + if tqdm is None: + tqdm = _tqdm + with tqdm(*args, **kwargs) as pbar: + with redirect_logging_to_tqdm(loggers=loggers, tqdm=tqdm): + yield pbar