Check column exists in NestedFrame.__setitem__

Co-authored-by: Wilson Beebe <[email protected]>
lincc-frameworks · Nov 6, 2024 · 1f86e14 · 1f86e14
1 parent 96503cd
commit 1f86e14
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git 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
@@ -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"]