Fix bug for NestedFrame.drop to support 'columns='

src/nested_pandas/nestedframe/core.py

@@ -777,10 +777,14 @@ def drop(
         """
 
         # axis 1 requires special handling for nested columns
-        if axis == 1:
+        if axis == 1 or columns is not None:
             # label convergence
             if isinstance(labels, str):
                 labels = [labels]
+            elif columns is not None:
+                labels = [columns] if isinstance(columns, str) else columns
+                columns = None
+                axis = 1
             nested_labels = [label for label in labels if self._is_known_hierarchical_column(label)]
             base_labels = [label for label in labels if not self._is_known_hierarchical_column(label)]
 

tests/nested_pandas/nestedframe/test_nestedframe.py

@@ -1434,6 +1434,15 @@ def test_drop():
     with pytest.raises(KeyError):
         base.drop(["a", "nested.not_a_field"], axis=1)
 
+    # Test dropping nested column using 'columns='
+    dropped_cols = base.drop(columns="nested.c")
+    assert "c" not in dropped_cols.nested.nest.columns
+
+    # Test dropping multiple nested columns using 'columns='
+    dropped_multcols = base.drop(columns=["nested.c", "nested2.f"])
+    assert "c" not in dropped_multcols.nested.nest.columns
+    assert "f" not in dropped_multcols.nested2.nest.columns
+
 
 def test_min():
     """Test min function return correct result with and without the nested columns"""