From e86b9b18a000136b85703b48504775d5aa350adc Mon Sep 17 00:00:00 2001
From: Daniel Ecer <de-code@users.noreply.github.com>
Date: Thu, 1 Apr 2021 20:32:38 +0100
Subject: [PATCH] added logging sub-module

added logging sub-module

python 2 compatibility

fixed python 2 fix

added test for custom tqdm class

python 2 absolute imports

(due to otherwise conflicting `logging` module)

isort

more tests relating to _get_first_found_console_formatter

isort

minor simplification

test handleError

test logging formatter being used

minor rename to _get_first_found_console_logging_formatter

test that certain exceptions are not swallowed

avoid using mock.assert_called (py 3.5)

moved to tqdm.contrib.logging

added "Redirecting console logging to tqdm" readme

removed no longer necessary absolute_import declaration

minor: updated package of example in docstring
---
 README.rst                     |  53 +++++++++
 tests/contrib/__init__.py      |   0
 tests/contrib/tests_logging.py | 182 +++++++++++++++++++++++++++++++
 tqdm/contrib/logging.py        | 194 +++++++++++++++++++++++++++++++++
 4 files changed, 429 insertions(+)
 create mode 100644 tests/contrib/__init__.py
 create mode 100644 tests/contrib/tests_logging.py
 create mode 100644 tqdm/contrib/logging.py

diff --git a/README.rst b/README.rst
index dd7e5a38..0ea0e98b 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 00000000..e69de29b
diff --git a/tests/contrib/tests_logging.py b/tests/contrib/tests_logging.py
new file mode 100644
index 00000000..0b334225
--- /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 00000000..08c9c3c4
--- /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