diff --git a/src/ape/logging.py b/src/ape/logging.py index fd4161c90d..33c6bc660f 100644 --- a/src/ape/logging.py +++ b/src/ape/logging.py @@ -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 @@ -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. diff --git a/tests/functional/test_logging.py b/tests/functional/test_logging.py index 9bf53cd768..b0c98bfb0f 100644 --- a/tests/functional/test_logging.py +++ b/tests/functional/test_logging.py @@ -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