Updated test cases to include more subplot stacking possibilities

Author: eicchen

pandas/tests/plotting/test_common.py

@@ -1,7 +1,6 @@
 import pytest
 import numpy as np
 from pandas import DataFrame
-from pandas import unique
 from pandas.tests.plotting.common import (
     _check_plot_works,
     _check_ticks_props,
@@ -59,45 +58,6 @@ def test_colorbar_layout(self):
 
         fig.colorbar(cs0, ax=[axes["A"], axes["B"]], location="right")
         DataFrame(x).plot(ax=axes["C"])
-    
-    def test_bar_subplot_stacking(self):
-        #GH Issue 61018
-        test_data = np.random.default_rng(3).integers(0,100,5)
-        df = DataFrame({"A": test_data, "B": test_data[::-1], "C": test_data[0]})
-        ax = df.plot(subplots= [('A','B')], kind="bar", stacked=True)
-
-        #finds all the rectangles that represent the values from both subplots
-        data_from_subplots = [[(x.get_x(), x.get_y(), x.get_height()) for x in ax[i].findobj(plt.Rectangle) if x.get_height() in test_data] for i in range(0,2)]
-
-        #get xy and height of squares that represent the data graphed from the df
-        #we would expect the height value of A to be reflected in the Y coord of B in subplot 1
-        subplot_data_df_list = []
-        unique_x_loc_list = []
-        for i in range(0,len(data_from_subplots)):
-            subplot_data_df= DataFrame(data = data_from_subplots[i], columns = ["x_coord", "y_coord", "height"])
-            unique_x_loc = unique(subplot_data_df["x_coord"])
-
-            subplot_data_df_list.append(subplot_data_df)
-            unique_x_loc_list.append(unique_x_loc)
-
-        #Checks subplot 1
-        plot_A_df = subplot_data_df_list[0].iloc[:len(test_data)]
-        plot_B_df = subplot_data_df_list[0].iloc[len(test_data):].reset_index()
-        total_bar_height = plot_A_df["height"].add(plot_B_df["height"])
-        #check number of bars matches the number of data plotted
-        assert len(unique_x_loc_list[0]) == len(test_data)
-        #checks that the first set of bars are the correct height and that the second one starts at the top of the first, additional checks the combined height of the bars are correct
-        assert (plot_A_df["height"] == test_data).all()
-        assert (plot_B_df["y_coord"] == test_data).all()
-        assert (total_bar_height == test_data + test_data[::-1]).all()
-
-        #Checks subplot 2
-        plot_C_df = subplot_data_df_list[1].iloc[:len(test_data)]
-        #check number of bars matches the number of data plotted
-        assert len(unique_x_loc_list[1]) == len(test_data)
-        #checks that all the bars start at zero and are the correct height
-        assert (plot_C_df["height"] == test_data[0]).all()
-        assert (plot_C_df["y_coord"] == 0).all()
 
 
 

pandas/tests/plotting/test_misc.py

@@ -681,3 +681,84 @@ def test_bar_plt_xaxis_intervalrange(self):
             (a.get_text() == b.get_text())
             for a, b in zip(s.plot.bar().get_xticklabels(), expected)
         )
+
+@pytest.fixture(scope="class")
+def BSS_data() -> np.array:
+    yield np.random.default_rng(3).integers(0,100,5)
+
+@pytest.fixture(scope="class")
+def BSS_df(BSS_data) -> DataFrame:
+    BSS_df = DataFrame({"A": BSS_data, "B": BSS_data[::-1], "C": BSS_data[0], "D": BSS_data[-1]})
+    return BSS_df
+
+def _BSS_xyheight_from_ax_helper(BSS_data, ax, subplot_division):
+    subplot_data_df_list = []
+
+    # get xy and height of squares that represent the data graphed from the df, seperated by subplots
+    for i in range(len(subplot_division)):
+        subplot_data = np.array([(x.get_x(), x.get_y(), x.get_height()) for x in ax[i].findobj(plt.Rectangle) if x.get_height() in BSS_data])
+        subplot_data_df_list.append(DataFrame(data = subplot_data, columns = ["x_coord", "y_coord", "height"]))
+
+    return subplot_data_df_list
+
+def _BSS_subplot_checker(BSS_data, BSS_df, subplot_data_df, subplot_columns):
+    assert_flag = 0
+    subplot_sliced_by_source = [subplot_data_df.iloc[len(BSS_data) * i : len(BSS_data) * (i+1)].reset_index() for i in range(0, len(subplot_columns))]
+    expected_total_height = BSS_df.loc[:,subplot_columns].sum(axis=1)
+    
+    for i in range(len(subplot_columns)):
+        sliced_df = subplot_sliced_by_source[i]
+        if i == 0:
+            #Checks that the bar chart starts y=0
+            assert (sliced_df["y_coord"] == 0).all
+            height_iter = sliced_df["y_coord"].add(sliced_df["height"])
+        else:
+            height_iter = height_iter + sliced_df["height"]
+
+        if i+1 == len(subplot_columns):
+            #Checks final height matches what is expected
+            tm.assert_series_equal(height_iter, expected_total_height, check_names = False, check_dtype= False)
+            
+        else:
+            #Checks each preceding bar ends where the next one starts
+            next_start_coord = subplot_sliced_by_source[i+1]["y_coord"]
+            tm.assert_series_equal(height_iter, next_start_coord, check_names = False, check_dtype= False)
+
+class TestBarSubplotStacked:
+    #GH Issue 61018
+    def test_bar_1_subplot_1_double_stacked(self, BSS_data, BSS_df):
+        columns_used = ["A", "B"]
+        BSS_df_trimmed = BSS_df[columns_used]
+        subplot_division = [columns_used]
+        ax = BSS_df_trimmed.plot(subplots = subplot_division, kind="bar", stacked=True)
+        subplot_data_df_list = _BSS_xyheight_from_ax_helper(BSS_data, ax, subplot_division)
+        for i in range(len(subplot_data_df_list)):
+            _BSS_subplot_checker(BSS_data, BSS_df_trimmed, subplot_data_df_list[i], subplot_division[i])
+        plt.savefig("1s1d.png")
+
+    
+    def test_bar_2_subplot_1_double_stacked(self, BSS_data, BSS_df):
+        columns_used = ["A", "B", "C"]
+        BSS_df_trimmed = BSS_df[columns_used]
+        subplot_division = [(columns_used[0], columns_used[1]), (columns_used[2],)]
+        ax = BSS_df_trimmed.plot(subplots = subplot_division, kind="bar", stacked=True)
+        subplot_data_df_list = _BSS_xyheight_from_ax_helper(BSS_data, ax, subplot_division)
+        for i in range(len(subplot_data_df_list)):
+            _BSS_subplot_checker(BSS_data, BSS_df_trimmed, subplot_data_df_list[i], subplot_division[i])
+        plt.savefig("2s1d.png")
+
+    def test_bar_2_subplot_2_double_stacked(self, BSS_data, BSS_df):
+        subplot_division = [('A', 'D'), ('C', 'B')]
+        ax = BSS_df.plot(subplots = subplot_division, kind="bar", stacked=True)
+        subplot_data_df_list = _BSS_xyheight_from_ax_helper(BSS_data, ax, subplot_division)
+        for i in range(len(subplot_data_df_list)):
+            _BSS_subplot_checker(BSS_data, BSS_df, subplot_data_df_list[i], subplot_division[i])
+        plt.savefig("2s2d.png")
+
+    def test_bar_2_subplots_1_triple_stacked(self, BSS_data, BSS_df):
+        subplot_division = [('A', 'D', 'C')]
+        ax = BSS_df.plot(subplots = subplot_division, kind="bar", stacked=True)
+        subplot_data_df_list = _BSS_xyheight_from_ax_helper(BSS_data, ax, subplot_division)
+        for i in range(len(subplot_data_df_list)):
+            _BSS_subplot_checker(BSS_data, BSS_df, subplot_data_df_list[i], subplot_division[i])
+        plt.savefig("2s1t.png")