FIX: match signature of axes.legend with add_sorted_driver_legend (#621)

---------

Co-authored-by: theOehrly <[email protected]>

fastf1/plotting/_interface.py

@@ -1,4 +1,5 @@
 import dataclasses
+import warnings
 from collections.abc import Sequence
 from typing import (
     Any,
@@ -8,6 +9,7 @@
 )
 
 import matplotlib.axes
+import matplotlib.legend
 
 from fastf1.core import Session
 from fastf1.internals.fuzzy import fuzzy_matcher
@@ -715,7 +717,9 @@ def list_compounds(session: Session) -> list[str]:
     return list(_Constants[year].CompoundColors.keys())
 
 
-def add_sorted_driver_legend(ax: matplotlib.axes.Axes, session: Session):
+def add_sorted_driver_legend(
+    ax: matplotlib.axes.Axes, session: Session, *args, **kwargs
+):
     """
     Adds a legend to the axis where drivers are grouped by team and within each
     team they are shown in the same order that is used for selecting plot
@@ -725,12 +729,18 @@ def add_sorted_driver_legend(ax: matplotlib.axes.Axes, session: Session):
     ``ax.legend()`` method. It can only be used when driver names or driver
     abbreviations are used as labels for the legend.
 
+    This function supports the same ``*args`` and ``**kwargs`` as
+    Matplotlib's ``ax.legend()``, including the ``handles`` and ``labels``
+    arguments. Check the Matplotlib documentation for more information.
+
     There is no particular need to use this function except to make the
     legend more visually pleasing.
 
     Args:
         ax: An instance of a Matplotlib ``Axes`` object
         session: the session for which the data should be obtained
+        *args: Matplotlib legend args
+        **kwargs: Matplotlib legend kwargs
 
      Returns:
         ``matplotlib.legend.Legend``
@@ -740,7 +750,22 @@ def add_sorted_driver_legend(ax: matplotlib.axes.Axes, session: Session):
 
     """
     dtm = _get_driver_team_mapping(session)
-    handles, labels = ax.get_legend_handles_labels()
+
+    try:
+        ret = matplotlib.legend._parse_legend_args([ax], *args, **kwargs)
+        if len(ret) == 3:
+            handles, labels, kwargs = ret
+            extra_args = []
+        else:
+            handles, labels, extra_args, kwargs = ret
+
+    except AttributeError:
+        warnings.warn("Failed to parse optional legend arguments correctly.",
+                      UserWarning)
+        extra_args = []
+        kwargs.pop('handles', None)
+        kwargs.pop('labels', None)
+        handles, labels = ax.get_legend_handles_labels()
 
     teams_list = list(dtm.teams_by_normalized.values())
     driver_list = list(dtm.drivers_by_normalized.values())
@@ -769,7 +794,7 @@ def add_sorted_driver_legend(ax: matplotlib.axes.Axes, session: Session):
         handles_new.append(elem[2])
         labels_new.append(elem[3])
 
-    return ax.legend(handles_new, labels_new)
+    return ax.legend(handles_new, labels_new, *extra_args, **kwargs)
 
 
 def set_default_colormap(colormap: str):

fastf1/tests/test_plotting.py

@@ -452,3 +452,12 @@ def test_override_team_constants():
 
     if fastf1.plotting.get_team_name('Haas', session, short=True) != 'Haas':
         raise RuntimeError("Test cleanup failed!")
+
+
+def test_import_internal_mpl_lgend_arg_kwarg_parser():
+    # Import the module and just try to access the internal function to see
+    # if it is available. This is not a test of the function itself but rather
+    # serves as an early warning if the function is removed or renamed by
+    # causing a test failure.
+    import matplotlib.legend
+    _ = matplotlib.legend._parse_legend_args