Improve with_log_level

usadelndsoc/__init__.py

 # -*- coding:utf-8; eval: (blacken-mode) -*-
 
 import logging
+import contextlib
 
 
 logger = logging.getLogger(__name__)
@@ -35,23 +36,20 @@ def _init_logging():
     logging.config.dictConfig(cfg)
 
 
-def with_log_level(level):
-    import functools
+class with_log_level(contextlib.ContextDecorator):
+    def __init__(self, level, old_level=None):
+        self._level = level
+        self._old_level = old_level
 
-    def deco(func):
-        @functools.wraps(func)
-        def wrapper(*a, **kw):
-            old_level = int(logger.level)
-            logger.setLevel(level)
-            try:
-                return func(*a, **kw)
-            finally:
-                logger.setLevel(old_level)
+    def __enter__(self):
+        cm = with_log_level(self._level, int(logger.level))
+        logger.setLevel(cm._level)
+        return cm
 
-        return wrapper
-
-    return deco
+    def __exit__(self, *exc):
+        logger.setLevel(self._old_level)
+        return False
 
 
 _init_logging()
-del _init_logging, logging
+del _init_logging, logging, contextlib