BUG Fixes fetch_kddcup99 for return_X_y and as_frame (#19011)

thomasjpfan · web-flow · commit eedda9191398 · 2020-12-16T09:58:39.000+01:00
diff --git a/sklearn/datasets/_kddcup99.py b/sklearn/datasets/_kddcup99.py
@@ -199,15 +199,15 @@ def fetch_kddcup99(*, subset=None, data_home=None, shuffle=False,
     with open(join(module_path, 'descr', 'kddcup99.rst')) as rst_file:
         fdescr = rst_file.read()
 
-    if return_X_y:
-        return data, target
-
     frame = None
     if as_frame:
         frame, data, target = _convert_data_dataframe(
             "fetch_kddcup99", data, target, feature_names, target_names
         )
 
+    if return_X_y:
+        return data, target
+
     return Bunch(
         data=data,
         target=target,
diff --git a/sklearn/datasets/tests/test_common.py b/sklearn/datasets/tests/test_common.py
@@ -75,6 +75,14 @@ def check_as_frame(bunch, dataset_func,
     if expected_target_dtype is not None:
         assert np.all(frame_bunch.target.dtypes == expected_target_dtype)
 
+    # Test for return_X_y and as_frame=True
+    frame_X, frame_y = dataset_func(as_frame=True, return_X_y=True)
+    assert isinstance(frame_X, pd.DataFrame)
+    if frame_y.ndim > 1:
+        assert isinstance(frame_X, pd.DataFrame)
+    else:
+        assert isinstance(frame_y, pd.Series)
+
 
 def _skip_network_tests():
     return os.environ.get('SKLEARN_SKIP_NETWORK_TESTS', '1') == '1'
