Merge pull request #173 from lincc-frameworks/setitem-missed-column

Check column exists in NestedFrame.__setitem__
lincc-frameworks · Nov 6, 2024 · 402ab66 · 402ab66
2 parents 96503cd + 1f86e14
commit 402ab66
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"]