From 99df05f3c064e55d55c49e80953618b0623eed19 Mon Sep 17 00:00:00 2001
From: antazoey <jules@apeworx.io>
Date: Thu, 19 Sep 2024 15:13:33 -0500
Subject: [PATCH] feat: add logger.at_level context manager method (#2294)

---
 src/ape/logging.py               | 20 +++++++++++++++++++-
 tests/functional/test_logging.py |  9 +++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

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