Support scoped enable and disable context managers

README.md

@@ -185,6 +185,20 @@ ic| 3: 3
 `ic()` continues to return its arguments when disabled, of course; no existing
 code with `ic()` breaks.
 
+Both methods also work as context managers. They restore the previous
+enabled state when the block exits, including when an exception is raised.
+Nested blocks restore their enclosing block's state.
+
+```python
+with ic.disable():
+    ic('hidden')
+    with ic.enable():
+        ic('visible inside the nested block')
+    ic('hidden again')
+
+ic('back to the original state')
+```
+
 
 ### Import Tricks
 

icecream/icecream.py

@@ -302,6 +302,19 @@ def _(obj: str) -> str:
     return "'" + obj.replace('\\', '\\\\') + "'"
 
 
+class _EnabledContext:
+    def __init__(self, debugger: 'IceCreamDebugger', enabled: bool) -> None:
+        self.debugger = debugger
+        self.previous = debugger.enabled
+        debugger.enabled = enabled
+
+    def __enter__(self) -> None:
+        return None
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.debugger.enabled = self.previous
+
+
 class IceCreamDebugger:
     _pairDelimiter = ', '  # Used by the tests in tests/.
     lineWrapWidth = DEFAULT_LINE_WRAP_WIDTH
@@ -472,11 +485,11 @@ def _getContext(self, callFrame: FrameType) -> Tuple[str, int, str]:
         filepath = (realpath if self.contextAbsPath else basename)(frameInfo.filename)  # type: ignore[operator]
         return filepath, lineNumber, parentFunction
 
-    def enable(self) -> None:
-        self.enabled = True
+    def enable(self) -> _EnabledContext:
+        return _EnabledContext(self, True)
 
-    def disable(self) -> None:
-        self.enabled = False
+    def disable(self) -> _EnabledContext:
+        return _EnabledContext(self, False)
 
     def use_stdout(self) -> None:
         if self.noColor:

tests/test_icecream.py

@@ -391,6 +391,35 @@ def test_enable_disable(self):
         pairs = parse_output_into_pairs(out, err, 2)
         assert pairs == [[('a', '1')], [('c', '3')]]
 
+    def test_enable_disable_contexts_restore_previous_state(self):
+        debugger = icecream.IceCreamDebugger()
+        for initially_enabled in (False, True):
+            for method in (debugger.enable, debugger.disable):
+                with self.subTest(initially_enabled=initially_enabled,
+                                  method=method.__name__):
+                    debugger.enabled = initially_enabled
+                    with method():
+                        self.assertEqual(debugger.enabled, method == debugger.enable)
+                    self.assertEqual(debugger.enabled, initially_enabled)
+
+    def test_enable_disable_nested_contexts(self):
+        debugger = icecream.IceCreamDebugger()
+        with debugger.disable():
+            with debugger.disable():
+                self.assertFalse(debugger.enabled)
+            self.assertFalse(debugger.enabled)
+            with debugger.enable():
+                self.assertTrue(debugger.enabled)
+            self.assertFalse(debugger.enabled)
+        self.assertTrue(debugger.enabled)
+
+    def test_enable_disable_context_propagates_exception(self):
+        debugger = icecream.IceCreamDebugger()
+        with self.assertRaisesRegex(ValueError, 'test failure'):
+            with debugger.disable():
+                raise ValueError('test failure')
+        self.assertTrue(debugger.enabled)
+
     def test_arg_to_string_function(self):
         def hello(obj):
             return 'zwei'