diff --git a/burr/core/state.py b/burr/core/state.py index f86c1c786..d5fbcc31e 100644 --- a/burr/core/state.py +++ b/burr/core/state.py @@ -436,8 +436,10 @@ def wipe(self, delete: Optional[list[str]] = None, keep: Optional[list[str]] = N ) if delete is not None: fields_to_delete = delete - else: + elif keep is not None: fields_to_delete = [key for key in self._state if key not in keep] + else: + fields_to_delete = list(self._state) return self.apply_operation(DeleteField(fields_to_delete)) def merge(self, other: "State") -> "State[StateType]": diff --git a/tests/core/test_state.py b/tests/core/test_state.py index c31a32ead..82a28678b 100644 --- a/tests/core/test_state.py +++ b/tests/core/test_state.py @@ -121,6 +121,13 @@ def test_state_wipe_keep(): assert wiped.get_all() == {"foo": "bar"} +def test_state_wipe_all(): + # wipe() with neither delete nor keep should clear everything, as documented. + state = State({"foo": "bar", "baz": "qux"}) + wiped = state.wipe() + assert wiped.get_all() == {} + + def test_state_append_validate_failure(): state = State({"foo": "bar"}) with pytest.raises(ValueError, match="non-appendable"):