From 30c6467e7802b60230b75f4ddcbc634862fffd4a Mon Sep 17 00:00:00 2001 From: Agampreet Singh Date: Fri, 3 Oct 2025 11:50:42 +0530 Subject: [PATCH 1/3] fix: State.keys() now displays only keys without values (#409) --- burr/core/state.py | 52 +++++++++++++++++++++++++++++++++++++++- tests/core/test_state.py | 48 +++++++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 1 deletion(-) diff --git a/burr/core/state.py b/burr/core/state.py index a9a31ac1d..540abd493 100644 --- a/burr/core/state.py +++ b/burr/core/state.py @@ -22,7 +22,19 @@ import inspect import logging from functools import cached_property -from typing import Any, Callable, Dict, Generic, Iterator, Mapping, Optional, TypeVar, Union +from typing import ( + Any, + Callable, + Dict, + Generic, + Iterator, + KeysView, + List, + Mapping, + Optional, + TypeVar, + Union, +) from burr.core import serde from burr.core.typing import DictBasedTypingSystem, TypingSystem @@ -270,6 +282,32 @@ def apply_mutate(self, inputs: dict): AssignedStateType = TypeVar("AssignedStateType") +class StateKeysView: + """Custom keys view that only displays keys, not values, for better usability. + + This addresses issue #409 by providing a clean representation of state keys + without showing the potentially large values that make debugging difficult. + """ + + def __init__(self, keys_view: KeysView[Any]) -> None: + self._keys_view = keys_view + + def __iter__(self) -> Iterator[Any]: + return iter(self._keys_view) + + def __len__(self) -> int: + return len(self._keys_view) + + def __contains__(self, key: Any) -> bool: + return key in self._keys_view + + def __repr__(self) -> str: + return f"StateKeys({list(self._keys_view)})" + + def __str__(self) -> str: + return repr(self) + + class State(Mapping, Generic[StateType]): """An immutable state object. This is the only way to interact with state in Burr.""" @@ -460,6 +498,18 @@ def __len__(self) -> int: def __iter__(self) -> Iterator[Any]: return iter(self._state) + def keys(self) -> StateKeysView: + """Returns a view of the state keys only (without values for cleaner display). + + This method addresses issue #409 by returning a custom view that displays + only the keys when printed, making it usable even when state contains + large objects like embeddings or dataframes. + + Returns: + StateKeysView: A view object that shows only keys in its string representation + """ + return StateKeysView(self._state.keys()) + def __repr__(self): return self.get_all().__repr__() # quick hack diff --git a/tests/core/test_state.py b/tests/core/test_state.py index dede0e99e..cccc42491 100644 --- a/tests/core/test_state.py +++ b/tests/core/test_state.py @@ -226,3 +226,51 @@ def test_state_apply_keeps_typing_system(): state = State({"foo": "bar"}, typing_system=SimpleTypingSystem()) assert state.update(foo="baz").typing_system is state.typing_system assert state.subset("foo").typing_system is state.typing_system + + +def test_state_keys_display(): + large_object = {"embedding": [0.1] * 100, "text": "large text" * 50} + state = State({"small_key": 42, "large_key": large_object}) + + keys = state.keys() + keys_repr = repr(keys) + + assert "small_key" in keys_repr + assert "large_key" in keys_repr + + assert "embedding" not in keys_repr + assert "large text" not in keys_repr + assert "0.1" not in keys_repr + + assert keys_repr.startswith("StateKeys(") + assert keys_repr.endswith("])") + + +def test_state_keys_functionality(): + state = State({"a": 1, "b": 2, "c": 3}) + keys = state.keys() + + assert list(keys) == ["a", "b", "c"] + + assert "a" in keys + assert "d" not in keys + + assert len(keys) == 3 + + empty_state = State() + empty_keys = empty_state.keys() + assert len(empty_keys) == 0 + assert list(empty_keys) == [] + + +def test_state_keys_compatibility(): + state = State({"x": 10, "y": 20}) + keys = state.keys() + + keys_set = set(keys) + assert keys_set == {"x", "y"} + + collected = [] + for key in keys: + collected.append(key) + assert collected == ["x", "y"] From cc1abc70190ac0b57855ee4da1109bacaf402bd9 Mon Sep 17 00:00:00 2001 From: Agampreet Singh Date: Sat, 4 Oct 2025 20:23:51 +0530 Subject: [PATCH 2/3] fix: return list from State.keys() to avoid displaying values (#409) --- burr/core/state.py | 39 ++++---------------------------- tests/core/test_state.py | 49 ++++++++++++++++++++++++++++++++++++++-- 2 files changed, 51 insertions(+), 37 deletions(-) diff --git a/burr/core/state.py b/burr/core/state.py index 540abd493..d82e6f4b7 100644 --- a/burr/core/state.py +++ b/burr/core/state.py @@ -28,7 +28,6 @@ Dict, Generic, Iterator, - KeysView, List, Mapping, Optional, @@ -282,32 +281,6 @@ def apply_mutate(self, inputs: dict): AssignedStateType = TypeVar("AssignedStateType") -class StateKeysView: - """Custom keys view that only displays keys, not values, for better usability. - - This addresses issue #409 by providing a clean representation of state keys - without showing the potentially large values that make debugging difficult. - """ - - def __init__(self, keys_view: KeysView[Any]) -> None: - self._keys_view = keys_view - - def __iter__(self) -> Iterator[Any]: - return iter(self._keys_view) - - def __len__(self) -> int: - return len(self._keys_view) - - def __contains__(self, key: Any) -> bool: - return key in self._keys_view - - def __repr__(self) -> str: - return f"StateKeys({list(self._keys_view)})" - - def __str__(self) -> str: - return repr(self) - - class State(Mapping, Generic[StateType]): """An immutable state object. This is the only way to interact with state in Burr.""" @@ -498,17 +471,13 @@ def __len__(self) -> int: def __iter__(self) -> Iterator[Any]: return iter(self._state) - def keys(self) -> StateKeysView: - """Returns a view of the state keys only (without values for cleaner display). - - This method addresses issue #409 by returning a custom view that displays - only the keys when printed, making it usable even when state contains - large objects like embeddings or dataframes. + def keys(self): + """Returns a list of the state keys only (without values for cleaner display). Returns: - StateKeysView: A view object that shows only keys in its string representation + list: A list of state keys """ - return StateKeysView(self._state.keys()) + return list(self._state) def __repr__(self): return self.get_all().__repr__() # quick hack diff --git a/tests/core/test_state.py b/tests/core/test_state.py index cccc42491..0521e388a 100644 --- a/tests/core/test_state.py +++ b/tests/core/test_state.py @@ -229,21 +229,66 @@ def test_state_apply_keeps_typing_system(): def test_state_keys_display(): + """Test that State.keys() displays only keys, not values (issue #409)""" + # Test with large objects that would make output unreadable large_object = {"embedding": [0.1] * 100, "text": "large text" * 50} state = State({"small_key": 42, "large_key": large_object}) keys = state.keys() keys_repr = repr(keys) + # Should return a simple list + assert isinstance(keys, list) + + # Should display only key names, not values assert "small_key" in keys_repr assert "large_key" in keys_repr + # Should NOT contain large object data assert "embedding" not in keys_repr assert "large text" not in keys_repr assert "0.1" not in keys_repr + + +def test_state_keys_functionality(): + """Test that State.keys() maintains expected list functionality""" + state = State({"a": 1, "b": 2, "c": 3}) + keys = state.keys() + + # Should return a list + assert isinstance(keys, list) + assert keys == ["a", "b", "c"] - assert keys_repr.startswith("StateKeys(") - assert keys_repr.endswith("])") + # Test list operations + assert "a" in keys + assert "d" not in keys + assert len(keys) == 3 + + # Test with empty state + empty_state = State() + empty_keys = empty_state.keys() + assert isinstance(empty_keys, list) + assert len(empty_keys) == 0 + assert empty_keys == [] + + +def test_state_keys_compatibility(): + """Test that State.keys() works with standard list operations""" + state = State({"x": 10, "y": 20}) + keys = state.keys() + + # Should work with list operations + assert isinstance(keys, list) + assert sorted(keys) == ["x", "y"] + + # Should work with indexing + assert keys[0] in ["x", "y"] + + # Should work in for loops + collected = [] + for key in keys: + collected.append(key) + assert set(collected) == {"x", "y"} def test_state_keys_functionality(): From d83e2919a367a02a80b5d351237d625930b79291 Mon Sep 17 00:00:00 2001 From: Agampreet Singh Date: Sat, 4 Oct 2025 21:27:37 +0530 Subject: [PATCH 3/3] removed redundant tests --- tests/core/test_state.py | 87 ++-------------------------------------- 1 file changed, 4 insertions(+), 83 deletions(-) diff --git a/tests/core/test_state.py b/tests/core/test_state.py index 0521e388a..1c19b8571 100644 --- a/tests/core/test_state.py +++ b/tests/core/test_state.py @@ -228,94 +228,15 @@ def test_state_apply_keeps_typing_system(): assert state.subset("foo").typing_system is state.typing_system -def test_state_keys_display(): - """Test that State.keys() displays only keys, not values (issue #409)""" - # Test with large objects that would make output unreadable - large_object = {"embedding": [0.1] * 100, "text": "large text" * 50} - state = State({"small_key": 42, "large_key": large_object}) - - keys = state.keys() - keys_repr = repr(keys) - - # Should return a simple list - assert isinstance(keys, list) - - # Should display only key names, not values - assert "small_key" in keys_repr - assert "large_key" in keys_repr - - # Should NOT contain large object data - assert "embedding" not in keys_repr - assert "large text" not in keys_repr - assert "0.1" not in keys_repr - - -def test_state_keys_functionality(): - """Test that State.keys() maintains expected list functionality""" +def test_state_keys_returns_list(): + """Test that State.keys() returns a list (fixes #409)""" state = State({"a": 1, "b": 2, "c": 3}) keys = state.keys() - # Should return a list + # Should return a list with the correct keys assert isinstance(keys, list) assert keys == ["a", "b", "c"] - # Test list operations - assert "a" in keys - assert "d" not in keys - assert len(keys) == 3 - # Test with empty state empty_state = State() - empty_keys = empty_state.keys() - assert isinstance(empty_keys, list) - assert len(empty_keys) == 0 - assert empty_keys == [] - - -def test_state_keys_compatibility(): - """Test that State.keys() works with standard list operations""" - state = State({"x": 10, "y": 20}) - keys = state.keys() - - # Should work with list operations - assert isinstance(keys, list) - assert sorted(keys) == ["x", "y"] - - # Should work with indexing - assert keys[0] in ["x", "y"] - - # Should work in for loops - collected = [] - for key in keys: - collected.append(key) - assert set(collected) == {"x", "y"} - - -def test_state_keys_functionality(): - state = State({"a": 1, "b": 2, "c": 3}) - keys = state.keys() - - assert list(keys) == ["a", "b", "c"] - - assert "a" in keys - assert "d" not in keys - - assert len(keys) == 3 - - empty_state = State() - empty_keys = empty_state.keys() - assert len(empty_keys) == 0 - assert list(empty_keys) == [] - - -def test_state_keys_compatibility(): - state = State({"x": 10, "y": 20}) - keys = state.keys() - - keys_set = set(keys) - assert keys_set == {"x", "y"} - - collected = [] - for key in keys: - collected.append(key) - assert collected == ["x", "y"] + assert empty_state.keys() == []