feat: add function to merge handles for identical labels in legend (#516

)

* Add function to merge handles for identical labels in legend

This function is useful, when using multiple plotting functions and one only wants to use a single label.
This function combines all the handles for identical labels into a single handle.

* Make merge_legend_handles_labels function available

src/mplhep/__init__.py

@@ -18,6 +18,7 @@
  hist2dplot,
  histplot,
  make_square_add_cbar,
+ merge_legend_handles_labels,
  mpl_magic,
  rescale_to_axessize,
  sort_legend,
@@ -70,6 +71,7 @@
  "rescale_to_axessize",
  "box_aspect",
  "make_square_add_cbar",
+ "merge_legend_handles_labels",
  "append_axes",
  "sort_legend",
  "save_variations",

src/mplhep/plot.py

@@ -1408,3 +1408,28 @@ def sort_legend(ax, order=None):
  if isinstance(order, OrderedDict):
  ordered_label_list = [order[k] for k in ordered_label_list]
  return ordered_label_values, ordered_label_list
+
+
+def merge_legend_handles_labels(handles, labels):
+ """
+ Merge handles for identical labels.
+ This is useful when combining multiple plot functions into a single label.
+
+ handles : List of handles
+ labels : List of labels
+ """
+
+ seen_labels = []
+ seen_label_handles = []
+ for handle, label in zip(handles, labels):
+ if label not in seen_labels:
+ seen_labels.append(label)
+ seen_label_handles.append([handle])
+ else:
+ idx = seen_labels.index(label)
+ seen_label_handles[idx].append(handle)
+
+ for i in range(len(seen_labels)):
+ seen_label_handles[i] = tuple(seen_label_handles[i])
+
+ return seen_label_handles, seen_labels