Preserve functionality for return_df arg

Author: nworb-cire

boruta/boruta_py.py

@@ -15,6 +15,7 @@
 from sklearn.base import BaseEstimator
 from sklearn.feature_selection import SelectorMixin
 from sklearn.utils.validation import check_is_fitted
+from sklearn.utils._set_output import _get_output_config
 import warnings
 
 
@@ -240,7 +241,7 @@ def transform(self, X, weak=None, return_df=None):
         weak : boolean, optional
             Deprecated. Set ``weak`` in the constructor instead.
 
-        return_df : bool, optional
+        return_df : boolean, optional
             Deprecated. Output type now follows scikit-learn's standard
             ``set_output``/``set_config`` mechanism.
         """
@@ -253,6 +254,9 @@ def transform(self, X, weak=None, return_df=None):
                 stacklevel=2,
             )
             self.weak = weak
+        requested_transform = None
+        prev_output_config = None
+        force_numpy = return_df is False
         if return_df is not None:
             warnings.warn(
                 "`return_df` is deprecated and will be removed in a future "
@@ -261,11 +265,20 @@ def transform(self, X, weak=None, return_df=None):
                 FutureWarning,
                 stacklevel=2,
             )
+            prev_output_config = _get_output_config("transform", estimator=self)["dense"]
+            requested_transform = "pandas" if return_df else "default"
+            if prev_output_config != requested_transform:
+                self.set_output(transform=requested_transform)
         try:
-            return super().transform(X)
+            result = super().transform(X)
         finally:
             if weak is not None:
                 self.weak = prev_weak
+            if requested_transform is not None and prev_output_config != requested_transform:
+                self.set_output(transform=prev_output_config)
+        if force_numpy and hasattr(result, "to_numpy"):
+            result = result.to_numpy()
+        return result
 
     def fit_transform(self, X, y=None, **fit_params):
         """

boruta/test/test_boruta.py

@@ -85,6 +85,33 @@ def test_return_df_parameter_emits_warning(Xy):
     assert isinstance(transformed, pd.DataFrame)
 
 
+def test_return_df_true_temporarily_enables_pandas_output(Xy):
+    X, y = Xy
+    bt = BorutaPy(RandomForestClassifier())
+    bt.fit(X, y)
+
+    baseline = bt.transform(X)
+    assert isinstance(baseline, np.ndarray)
+
+    with pytest.warns(FutureWarning, match="`return_df` is deprecated"):
+        transformed = bt.transform(X, return_df=True)
+    assert isinstance(transformed, pd.DataFrame)
+
+    reverted = bt.transform(X)
+    assert isinstance(reverted, np.ndarray)
+
+
+def test_return_df_false_with_dataframe_input_returns_numpy(Xy):
+    X, y = Xy
+    X_df = pd.DataFrame(X)
+    bt = BorutaPy(RandomForestClassifier())
+    bt.fit(X_df, y)
+
+    with pytest.warns(FutureWarning, match="`return_df` is deprecated"):
+        transformed = bt.transform(X_df, return_df=False)
+    assert isinstance(transformed, np.ndarray)
+
+
 def test_weak_attribute_controls_support_mask(Xy):
     X, y = Xy
     bt = BorutaPy(RandomForestClassifier(), weak=True)