Skip to content

Commit

Permalink
feat: add logger.at_level context manager method (#2294)
Browse files Browse the repository at this point in the history
  • Loading branch information
antazoey committed Sep 19, 2024
1 parent 07a792c commit 99df05f
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 1 deletion.
20 changes: 19 additions & 1 deletion src/ape/logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@
import logging
import sys
import traceback
from collections.abc import Callable, Sequence
from collections.abc import Callable, Iterator, Sequence
from contextlib import contextmanager
from enum import IntEnum
from pathlib import Path
from typing import IO, Any, Optional, Union
Expand Down Expand Up @@ -185,6 +186,23 @@ def set_level(self, level: Union[str, int, LogLevel]):
for _logger in self._extra_loggers.values():
_logger.setLevel(level)

@contextmanager
def at_level(self, level: Union[str, int, LogLevel]) -> Iterator:
"""
Change the log-level in a context.
Args:
level (Union[str, int, LogLevel]): The level to use.
Returns:
Iterator
"""

initial_level = self.level
self.set_level(level)
yield
self.set_level(initial_level)

def log_error(self, err: Exception):
"""
Avoids logging empty messages.
Expand Down
9 changes: 9 additions & 0 deletions tests/functional/test_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,3 +107,12 @@ def cmd(cli_ctx):
def test_set_level(level):
logger.set_level(level)
assert logger.level == LogLevel.INFO.value


def test_at_level():
level_to_set = next(lvl for lvl in LogLevel if lvl != logger.level)
initial_level = logger.level
with logger.at_level(level_to_set):
assert logger.level == level_to_set

assert logger.level == initial_level

0 comments on commit 99df05f

Please sign in to comment.