diff --git a/taskiq/state.py b/taskiq/state.py index 1629aecd..40dfc672 100644 --- a/taskiq/state.py +++ b/taskiq/state.py @@ -41,3 +41,7 @@ def __delattr__(self, name: str) -> None: def __str__(self) -> str: return f"TaskiqState({super().__str__()})" + + def copy(self) -> "TaskiqState": + """Return a shallow copy without changing the original state.""" + return self.__copy__() diff --git a/tests/test_state.py b/tests/test_state.py index 09f1efaf..f6ac19a8 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -56,3 +56,28 @@ def test_state_del_attr() -> None: del state.a assert state.get("a") is None + + +def test_state_copy() -> None: + state = TaskiqState() + values = [1] + state.items_list = values + + copied = state.copy() + + assert isinstance(copied, TaskiqState) + assert dict(state) == dict(copied) == {"items_list": values} + assert copied["items_list"] is values + copied.other = 2 + assert "other" not in state + + +def test_state_copy_preserves_data_key() -> None: + state = TaskiqState() + state["data"] = "user value" + + copied = state.copy() + + assert dict(state) == dict(copied) == {"data": "user value"} + copied["data"] = "copied value" + assert state["data"] == "user value"