diff --git a/cid/set.py b/cid/set.py index ab5d4e1..3f28eb8 100644 --- a/cid/set.py +++ b/cid/set.py @@ -1,7 +1,7 @@ """CID Set operations for managing collections of unique CIDs.""" from collections.abc import Callable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from .cid import CIDv0, CIDv1 @@ -77,15 +77,22 @@ def visit(self, cid: "CIDv0 | CIDv1") -> bool: return True return False - def for_each(self, func: Callable[["CIDv0 | CIDv1"], None]) -> None: + def for_each(self, func: Callable[["CIDv0 | CIDv1"], "Any"]) -> "Exception | None": """ - Call function for each CID in set. + Call function for each CID in set. Stop and return error if func raises or returns an error. :param func: Function to call for each CID :type func: callable + :return: Exception if an error occurred, None otherwise """ for cid in self._set: - func(cid) + try: + result = func(cid) + if result is not None and isinstance(result, Exception): + return result + except Exception as e: + return e + return None def __iter__(self): """Make set iterable.""" diff --git a/newsfragments/69.bugfix.rst b/newsfragments/69.bugfix.rst new file mode 100644 index 0000000..9e6e953 --- /dev/null +++ b/newsfragments/69.bugfix.rst @@ -0,0 +1 @@ +Update `CIDSet.for_each()` to stop iteration and return the error if the callback function raises an exception or returns an exception. diff --git a/tests/test_new_features.py b/tests/test_new_features.py index 04c766e..afb570d 100644 --- a/tests/test_new_features.py +++ b/tests/test_new_features.py @@ -262,6 +262,26 @@ def collect(cid): assert cidv0 in collected assert cidv1 in collected + def test_cid_set_for_each_error_propagation(self, cidv0, cidv1): + """CIDSet.for_each: stops and propagates error""" + cid_set = CIDSet() + cid_set.add(cidv0) + cid_set.add(cidv1) + + def failing_func(cid): + raise ValueError("test error") + + result = cid_set.for_each(failing_func) + assert isinstance(result, ValueError) + assert str(result) == "test error" + + def returning_func(cid): + return TypeError("return error") + + result2 = cid_set.for_each(returning_func) + assert isinstance(result2, TypeError) + assert str(result2) == "return error" + def test_cid_set_contains(self, cidv0): """CIDSet.__contains__: supports 'in' operator""" cid_set = CIDSet()