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

Check column exists in NestedFrame.__setitem__

Author: wilsonbb

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)
 

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"]