mirror of https://github.com/tqdm/tqdm.git
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:
parent
46907ea524
commit
f99bcb479b
|
@ -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:
|
||||||
|
|
|
@ -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]
|
||||||
|
|
Loading…
Reference in New Issue