Skip to content

Commit

Permalink
added logging sub-module
Browse files Browse the repository at this point in the history
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)
  • Loading branch information
de-code committed Apr 2, 2021
1 parent bcce20f commit c24dab3
Show file tree
Hide file tree
Showing 3 changed files with 378 additions and 0 deletions.
182 changes: 182 additions & 0 deletions tests/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.logging import _get_first_found_console_logging_formatter
from tqdm.logging import _TqdmLoggingHandler as TqdmLoggingHandler
from tqdm.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']
2 changes: 2 additions & 0 deletions tqdm/cli.py
@@ -1,6 +1,8 @@
"""
Module version for monitoring CLI pipes (`... | python -m tqdm | ...`).
"""
from __future__ import absolute_import

import logging
import re
import sys
Expand Down
194 changes: 194 additions & 0 deletions tqdm/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.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.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.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

0 comments on commit c24dab3

Please sign in to comment.