Skip to content

feat(State): Implement __delitem__ and pop methods for element removal #1856

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions src/google/adk/sessions/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@

from typing import Any

_sentinel = object()


class State:
"""A state dict that maintain the current value and the pending-commit delta."""
Expand All @@ -37,6 +39,15 @@ def __getitem__(self, key: str) -> Any:
return self._delta[key]
return self._value[key]

def __delitem__(self, key: str):
"""Deletes the value of the state dict for the given key"""
if key not in self:
raise KeyError(key)
if key in self._delta:
del self._delta[key]
if key in self._value:
del self._value[key]

def __setitem__(self, key: str, value: Any):
"""Sets the value of the state dict for the given key."""
# TODO: make new change only store in delta, so that self._value is only
Expand All @@ -58,6 +69,18 @@ def get(self, key: str, default: Any = None) -> Any:
return default
return self[key]

def pop(self, key: str, default: Any = _sentinel) -> Any:
"""Deletes the value of the state dict for the given key"""
if key in self:
value_to_return = self[key]
del self[key]
return value_to_return
else:
if default is not _sentinel:
return default
else:
raise KeyError(key)

def update(self, delta: dict[str, Any]):
"""Updates the state dict with the given delta."""
self._value.update(delta)
Expand Down