chore: move legend logic to standard graph utils

Author: Mrglglglglgl

pyretailscience/plots/line.py

@@ -45,14 +45,14 @@
 def plot(
     df: pd.DataFrame,
     value_col: str | list[str],
-    x_label: str | None = None,
-    y_label: str | None = None,
-    title: str | None = None,
+    x_label: str,
+    y_label: str,
+    title: str,
     x_col: str | None = None,
     group_col: str | None = None,
-    legend_title: str | None = None,
     ax: Axes | None = None,
     source_text: str | None = None,
+    legend_title: str | None = None,
     move_legend_outside: bool = False,
     **kwargs: dict[str, any],
 ) -> SubplotBase:
@@ -61,9 +61,9 @@ def plot(
     Args:
         df (pd.DataFrame): The dataframe to plot.
         value_col (str or list of str): The column(s) to plot.
-        x_label (str, optional): The x-axis label.
-        y_label (str, optional): The y-axis label.
-        title (str, optional): The title of the plot.
+        x_label (str): The x-axis label.
+        y_label (str): The y-axis label.
+        title (str): The title of the plot.
         x_col (str, optional): The column to be used as the x-axis. If None, the index is used.
         group_col (str, optional): The column used to define different lines.
         legend_title (str, optional): The title of the legend.
@@ -96,15 +96,14 @@ def plot(
         **kwargs,
     )
 
-    ax = gu.standard_graph_styles(ax=ax, title=title, x_label=x_label, y_label=y_label)
-
-    if move_legend_outside:
-        ax.legend(bbox_to_anchor=(1.05, 1))
-
-    if legend_title is not None:
-        legend = ax.get_legend()
-        if legend is not None:
-            legend.set_title(legend_title)
+    ax = gu.standard_graph_styles(
+        ax=ax,
+        title=title,
+        x_label=x_label,
+        y_label=y_label,
+        legend_title=legend_title,
+        move_legend_outside=move_legend_outside,
+    )
 
     if source_text is not None:
         gu.add_source_text(ax=ax, source_text=source_text)

pyretailscience/style/graph_utils.py

@@ -67,6 +67,8 @@ def standard_graph_styles(
     title_pad: int = GraphStyles.DEFAULT_TITLE_PAD,
     x_label_pad: int = GraphStyles.DEFAULT_AXIS_LABEL_PAD,
     y_label_pad: int = GraphStyles.DEFAULT_AXIS_LABEL_PAD,
+    legend_title: str | None = None,
+    move_legend_outside: bool = False,
 ) -> Axes:
     """Apply standard styles to a Matplotlib graph.
 
@@ -79,6 +81,8 @@ def standard_graph_styles(
         x_label_pad (int, optional): The padding below the x-axis label. Defaults to GraphStyles.DEFAULT_AXIS_LABEL_PAD.
         y_label_pad (int, optional): The padding to the left of the y-axis label. Defaults to
             GraphStyles.DEFAULT_AXIS_LABEL_PAD.
+        legend_title (str, optional): The title of the legend. If None, no legend title is applied. Defaults to None.
+        move_legend_outside (bool, optional): Whether to move the legend outside the plot. Defaults to False.
 
     Returns:
         Axes: The graph with the styles applied.
@@ -112,6 +116,13 @@ def standard_graph_styles(
             labelpad=y_label_pad,
         )
 
+    legend = ax.legend() if move_legend_outside or legend_title is not None else ax.get_legend()
+    if legend:
+        if move_legend_outside:
+            legend.set_bbox_to_anchor((1.05, 1))
+        if legend_title:
+            legend.set_title(legend_title)
+
     return ax
 
 

tests/plots/test_line.py

@@ -107,7 +107,7 @@ def test_plot_moves_legend_outside(sample_dataframe):
     """Test the plot function moves the legend outside the plot."""
     _, ax = plt.subplots()
 
-    # Test with move_legend_outside=True
+    # Create the plot with move_legend_outside=True
     result_ax = line.plot(
         df=sample_dataframe,
         value_col="y",
@@ -120,13 +120,44 @@ def test_plot_moves_legend_outside(sample_dataframe):
         move_legend_outside=True,
     )
 
-    expected_coords = (1.05, 1.0)
-    legend = result_ax.get_legend()
-    # Check if bbox_to_anchor is set to (1.05, 1) when legend is outside
-    bbox_anchor = legend.get_bbox_to_anchor()._bbox
+    # Assert that standard_graph_styles was called with move_legend_outside=True
+    gu.standard_graph_styles.assert_called_once_with(
+        ax=result_ax,
+        title="Test Plot Legend Outside",
+        x_label="X Axis",
+        y_label="Y Axis",
+        legend_title=None,
+        move_legend_outside=True,
+    )
+
+
+@pytest.mark.usefixtures("_mock_get_base_cmap", "_mock_gu_functions")
+def test_plot_moves_legend_inside(sample_dataframe):
+    """Test the plot function moves the legend inside the plot."""
+    _, ax = plt.subplots()
+
+    # Create the plot with move_legend_outside=False
+    result_ax = line.plot(
+        df=sample_dataframe,
+        value_col="y",
+        x_label="X Axis",
+        y_label="Y Axis",
+        title="Test Plot Legend Inside",
+        x_col="x",
+        group_col="group",
+        ax=ax,
+        move_legend_outside=False,
+    )
 
-    assert legend is not None
-    assert (bbox_anchor.x0, bbox_anchor.y0) == expected_coords
+    # Assert that standard_graph_styles was called with move_legend_outside=False
+    gu.standard_graph_styles.assert_called_once_with(
+        ax=result_ax,
+        title="Test Plot Legend Inside",
+        x_label="X Axis",
+        y_label="Y Axis",
+        legend_title=None,
+        move_legend_outside=False,
+    )
 
 
 @pytest.mark.usefixtures("_mock_get_base_cmap", "_mock_gu_functions")