From 1f86e14cbe8ab7edb785eb49cc1ead580643bfd9 Mon Sep 17 00:00:00 2001 From: Konstantin Malanchev Date: Wed, 6 Nov 2024 17:24:40 -0500 Subject: [PATCH] Check column exists in NestedFrame.__setitem__ Co-authored-by: Wilson Beebe --- src/nested_pandas/nestedframe/core.py | 2 ++ tests/nested_pandas/nestedframe/test_nestedframe.py | 7 +++++++ 2 files changed, 9 insertions(+) diff --git a/src/nested_pandas/nestedframe/core.py b/src/nested_pandas/nestedframe/core.py index a0674b8..192a421 100644 --- a/src/nested_pandas/nestedframe/core.py +++ b/src/nested_pandas/nestedframe/core.py @@ -202,6 +202,8 @@ def __getitem__(self, item): nested = item.split(".")[0] col = ".".join(item.split(".")[1:]) return self[nested].nest.get_flat_series(col) + else: + raise KeyError(f"Column '{item}' not found in nested columns or base columns") else: return super().__getitem__(item) diff --git a/tests/nested_pandas/nestedframe/test_nestedframe.py b/tests/nested_pandas/nestedframe/test_nestedframe.py index 9175b3d..938476f 100644 --- a/tests/nested_pandas/nestedframe/test_nestedframe.py +++ b/tests/nested_pandas/nestedframe/test_nestedframe.py @@ -1014,3 +1014,10 @@ def test_eval_assignment(): assert set(nf.p2.nest.fields) == {"e", "f"} assert (nf["p2.e"] == nf["packed.d"] * 2 + nf.c).all() assert (nf["p2.f"] == nf["p2.e"] + nf.b).all() + + +def test_access_non_existing_column(): + """Test that accessing a non-existing column raises a KeyError""" + nf = NestedFrame() + with pytest.raises(KeyError): + _ = nf["non_existing_column"]