logging: inherit output stream from existing handler

When using the logging redirection, logs will currently always be printed
to stdout, while the logging module default is to print to stderr.

Fix this by trying to inherit the stream from the existing handler, like
the code already does for the formatter.

Consider the following example:

    import logging
    import time

    from tqdm.contrib.logging import tqdm_logging_redirect

    log = logging.getLogger()
    log.warning("start")

    with tqdm_logging_redirect(range(int(4))) as pbar:
        for i in pbar:
            time.sleep(0.1)
            log.warning(f"Step {i}")

    log.warning("done")

Running this while redirecting stdout (`$ python3 log.py > /dev/null`)
without this patch will print:

    $ venv/bin/python log.py > /dev/null
    start
    100%|████████████████████████████████████████████| 4/4 [00:00<00:00,  9.87it/s]
    done

After this patch:

    $ venv/bin/python log.py > /dev/null
    start
    Step 0
    Step 1
    Step 2
    Step 3
    100%|████████████████████████████████████████████| 4/4 [00:00<00:00,  9.83it/s]
    done

Signed-off-by: Steffan Karger <steffan.karger@fox-it.com>
This commit is contained in:
Steffan Karger 2021-06-21 11:50:26 +02:00 committed by Casper da Costa-Luis
parent 46907ea524
commit f99bcb479b
2 changed files with 16 additions and 22 deletions

View File

@ -10,7 +10,7 @@ from io import StringIO
import pytest import pytest
from tqdm import tqdm from tqdm import tqdm
from tqdm.contrib.logging import _get_first_found_console_logging_formatter from tqdm.contrib.logging import _get_first_found_console_logging_handler
from tqdm.contrib.logging import _TqdmLoggingHandler as TqdmLoggingHandler from tqdm.contrib.logging import _TqdmLoggingHandler as TqdmLoggingHandler
from tqdm.contrib.logging import logging_redirect_tqdm, tqdm_logging_redirect from tqdm.contrib.logging import logging_redirect_tqdm, tqdm_logging_redirect
@ -68,33 +68,25 @@ class TestTqdmLoggingHandler:
logger.info('test') logger.info('test')
class TestGetFirstFoundConsoleLoggingFormatter: class TestGetFirstFoundConsoleLoggingHandler:
def test_should_return_none_for_no_handlers(self): def test_should_return_none_for_no_handlers(self):
assert _get_first_found_console_logging_formatter([]) is None assert _get_first_found_console_logging_handler([]) is None
def test_should_return_none_without_stream_handler(self): def test_should_return_none_without_stream_handler(self):
handler = logging.handlers.MemoryHandler(capacity=1) handler = logging.handlers.MemoryHandler(capacity=1)
handler.formatter = TEST_LOGGING_FORMATTER assert _get_first_found_console_logging_handler([handler]) is None
assert _get_first_found_console_logging_formatter([handler]) is None
def test_should_return_none_for_stream_handler_not_stdout_or_stderr(self): def test_should_return_none_for_stream_handler_not_stdout_or_stderr(self):
handler = logging.StreamHandler(StringIO()) handler = logging.StreamHandler(StringIO())
handler.formatter = TEST_LOGGING_FORMATTER assert _get_first_found_console_logging_handler([handler]) is None
assert _get_first_found_console_logging_formatter([handler]) is None
def test_should_return_stream_handler_formatter_if_stream_is_stdout(self): def test_should_return_stream_handler_if_stream_is_stdout(self):
handler = logging.StreamHandler(sys.stdout) handler = logging.StreamHandler(sys.stdout)
handler.formatter = TEST_LOGGING_FORMATTER assert _get_first_found_console_logging_handler([handler]) == handler
assert _get_first_found_console_logging_formatter(
[handler]
) == TEST_LOGGING_FORMATTER
def test_should_return_stream_handler_formatter_if_stream_is_stderr(self): def test_should_return_stream_handler_if_stream_is_stderr(self):
handler = logging.StreamHandler(sys.stderr) handler = logging.StreamHandler(sys.stderr)
handler.formatter = TEST_LOGGING_FORMATTER assert _get_first_found_console_logging_handler([handler]) == handler
assert _get_first_found_console_logging_formatter(
[handler]
) == TEST_LOGGING_FORMATTER
class TestRedirectLoggingToTqdm: class TestRedirectLoggingToTqdm:

View File

@ -26,7 +26,7 @@ class _TqdmLoggingHandler(logging.StreamHandler):
def emit(self, record): def emit(self, record):
try: try:
msg = self.format(record) msg = self.format(record)
self.tqdm_class.write(msg) self.tqdm_class.write(msg, file=self.stream)
self.flush() self.flush()
except (KeyboardInterrupt, SystemExit): except (KeyboardInterrupt, SystemExit):
raise raise
@ -39,10 +39,10 @@ def _is_console_logging_handler(handler):
and handler.stream in {sys.stdout, sys.stderr}) and handler.stream in {sys.stdout, sys.stderr})
def _get_first_found_console_logging_formatter(handlers): def _get_first_found_console_logging_handler(handlers):
for handler in handlers: for handler in handlers:
if _is_console_logging_handler(handler): if _is_console_logging_handler(handler):
return handler.formatter return handler
@contextmanager @contextmanager
@ -85,8 +85,10 @@ def logging_redirect_tqdm(
try: try:
for logger in loggers: for logger in loggers:
tqdm_handler = _TqdmLoggingHandler(tqdm_class) tqdm_handler = _TqdmLoggingHandler(tqdm_class)
tqdm_handler.setFormatter( orig_handler = _get_first_found_console_logging_handler(logger.handlers)
_get_first_found_console_logging_formatter(logger.handlers)) if orig_handler is not None:
tqdm_handler.setFormatter(orig_handler.formatter)
tqdm_handler.stream = orig_handler.stream
logger.handlers = [ logger.handlers = [
handler for handler in logger.handlers handler for handler in logger.handlers
if not _is_console_logging_handler(handler)] + [tqdm_handler] if not _is_console_logging_handler(handler)] + [tqdm_handler]