Fix bug for NestedFrame.drop to support 'columns=' (#407)

* update drop support via 'columns=' and add tests

* delete comment

* remove redundant logic

* fix docstring link

Author: Graciaaa3

src/nested_pandas/nestedframe/core.py

@@ -727,7 +727,8 @@ def drop(
         Remove rows or columns by specifying label names and corresponding
         axis, or by directly specifying index or column names. When using a
         multi-index, labels on different levels can be removed by
-        specifying the level. See the user guide for more information about
+        specifying the level. See the `user guide <https://pandas.pydata.org/docs/user_guide
+        /advanced.html#advanced-shown-levels>`_ for more information about
         the now unused levels.
 
         Parameters
@@ -777,10 +778,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"""