Revert "Add GridSearchColorPlot tests + minor gridsearch change"

This reverts commit 5ebca26.
DistrictDataLabs · Oct 11, 2020 · 56de311 · 56de311
1 parent 5845189
commit 56de311
Show file tree

Hide file tree

Showing 6 changed files with 4 additions and 135 deletions.
diff --git a/tests/baseline_images/test_gridsearch/test_base/test_gridsearchcolorplot.png b/tests/baseline_images/test_gridsearch/test_base/test_gridsearchcolorplot.png
diff --git a/tests/baseline_images/test_gridsearch/test_base/test_numpy_integration.png b/tests/baseline_images/test_gridsearch/test_base/test_numpy_integration.png
diff --git a/tests/baseline_images/test_gridsearch/test_base/test_pandas_integration.png b/tests/baseline_images/test_gridsearch/test_base/test_pandas_integration.png
diff --git a/tests/baseline_images/test_gridsearch/test_base/test_quick_method.png b/tests/baseline_images/test_gridsearch/test_base/test_quick_method.png
diff --git a/tests/test_gridsearch/test_base.py b/tests/test_gridsearch/test_base.py
diff --git a/yellowbrick/gridsearch/pcolor.py b/yellowbrick/gridsearch/pcolor.py
@@ -68,8 +68,8 @@ def gridsearch_color_plot(model, x_param, y_param, X=None, y=None, ax=None, **kw
 
     Returns
     -------
-    visualizer : GridSearchColorPlot
-        Returns visualizer
+    ax : matplotlib axes
+        Returns the axes that the classification report was drawn on.
     """
     # Instantiate the visualizer
     visualizer = GridSearchColorPlot(model, x_param, y_param, ax=ax, **kwargs)
@@ -80,8 +80,8 @@ def gridsearch_color_plot(model, x_param, y_param, X=None, y=None, ax=None, **kw
     else:
         visualizer.draw()
 
-    # Return the visualizer
-    return visualizer
+    # Return the axes object on the visualizer
+    return visualizer.ax
 
 
 class GridSearchColorPlot(GridSearchVisualizer):