Add calling() method to CircuitBreaker, returning a context manager (#87

)

CHANGELOG

@@ -1,6 +1,10 @@
 Changelog
 =========
 
+Version 1.1.0 (January 09, 2024)
+
+* Added calling() method to CircuitBreaker, returning a context manager (Thanks @martijnthe)
+
 Version 1.0.3 (December 22, 2023)
 
 * Added py.typed file (Thanks @martijnthe)

README.rst

@@ -196,6 +196,16 @@ Or if you don't want to use the decorator syntax:
     # Will trigger the circuit breaker
     updated_customer = db_breaker.call(update_customer, my_customer)
 
+Or use it as a context manager and a `with` statement:
+
+.. code:: python
+
+    # Will trigger the circuit breaker
+    with db_breaker.calling():
+        # Do stuff here...
+        pass
+
+
 
 According to the default parameters, the circuit breaker ``db_breaker`` will
 automatically open the circuit after 5 consecutive failures in

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="pybreaker",
-    version="1.0.3",
+    version="1.1.0",
     description="Python implementation of the Circuit Breaker pattern",
     long_description=open("README.rst", "r").read(),
     keywords=["design", "pattern", "circuit", "breaker", "integration"],

src/pybreaker/__init__.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import calendar
+import contextlib
 import logging
 import sys
 import threading
@@ -27,6 +28,7 @@
     Union,
     cast,
     overload,
+    Iterable,
 )
 
 if sys.version_info >= (3, 8):
@@ -260,6 +262,19 @@ def call(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> T:
         with self._lock:
             return self.state.call(func, *args, **kwargs)
 
+    @contextlib.contextmanager
+    def calling(self) -> Iterable[None]:
+        """
+        Returns a context manager, enabling the circuit breaker to be used with a
+        `with` statement. The block of code inside the `with` statement will be
+        executed according to the rules implemented by the current state of this
+        circuit breaker.
+        """
+        def _wrapper() -> None:
+            yield
+
+        yield from self.call(_wrapper)
+
     def call_async(self, func, *args, **kwargs):  # type: ignore[no-untyped-def]
         """
         Calls async `func` with the given `args` and `kwargs` according to the rules

src/tests.py

@@ -1159,5 +1159,46 @@ def before_call(self, cb, func, *args, **kwargs):
         self.assertTrue(self.breaker.fail_counter < self.breaker.fail_max + num_threads)
 
 
+class CircuitBreakerContextManagerTestCase(unittest.TestCase):
+    """
+    Tests for the CircuitBreaker class, when used as a context manager.
+    """
+
+    def test_calling(self):
+        """
+        Test that the CircuitBreaker calling() API returns a context manager and works as expected.
+        """
+
+        class TestError(Exception):
+            pass
+
+        breaker = CircuitBreaker(fail_max=2, reset_timeout=0.01)
+        mock_fn = mock.MagicMock()
+
+        def _do_raise():
+            with breaker.calling():
+                raise TestError
+
+        def _do_succeed():
+            with breaker.calling():
+                mock_fn()
+
+        self.assertRaises(TestError, _do_raise)
+        self.assertRaises(CircuitBreakerError, _do_raise)
+        self.assertEqual(2, breaker.fail_counter)
+        self.assertEqual("open", breaker.current_state)
+
+        # Still fails while circuit breaker is open:
+        self.assertRaises(CircuitBreakerError, _do_succeed)
+        mock_fn.assert_not_called()
+
+        sleep(0.01)
+
+        _do_succeed()
+        mock_fn.assert_called_once()
+        self.assertEqual(0, breaker.fail_counter)
+        self.assertEqual("closed", breaker.current_state)
+
+
 if __name__ == "__main__":
     unittest.main()