Fix cphyc#87

Only label those lines that are passed to the labelLines function

Author: cphyc

labellines/core.py

@@ -1,5 +1,6 @@
 import warnings
 
+import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.container import ErrorbarContainer
 from matplotlib.dates import DateConverter, num2date
@@ -81,7 +82,7 @@ def labelLine(
 
 
 def labelLines(
-    lines,
+    lines=None,
     align=True,
     xvals=None,
     drop_label=False,
@@ -95,8 +96,8 @@ def labelLines(
 
     Parameters
     ----------
-    lines : list of matplotlib lines
-       The lines to label
+    lines : list of matplotlib lines, optional.
+       Lines to label. If empty, label all lines that have a label.
     align : boolean, optional
        If True, the label will be aligned with the slope of the line
        at the location of the label. If False, they will be horizontal.
@@ -119,16 +120,34 @@ def labelLines(
     kwargs : dict, optional
        Optional arguments passed to ax.text
     """
-    ax = lines[0].axes
+    if lines:
+        ax = lines[0].axes
+    else:
+        ax = plt.gca()
 
     handles, allLabels = ax.get_legend_handles_labels()
 
     all_lines = []
     for h in handles:
         if isinstance(h, ErrorbarContainer):
-            all_lines.append(h.lines[0])
+            line = h.lines[0]
         else:
-            all_lines.append(h)
+            line = h
+
+        if (lines is not None) and (line not in lines):
+            continue
+        all_lines.append(line)
+
+    # Check that the lines passed to the function have all a label
+    if lines is not None:
+        for line in lines:
+            if line in all_lines:
+                continue
+
+            warnings.warn(
+                "Tried to label line %s, but could not find a label for it.",
+                UserWarning,
+            )
 
     # In case no x location was provided, we need to use some heuristics
     # to generate them.

labellines/test.py

@@ -1,12 +1,10 @@
-import warnings
 from datetime import datetime
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pytest
 from matplotlib.dates import UTC, DateFormatter, DayLocator
 from matplotlib.testing import setup
-from numpy.testing import assert_raises
 
 from .core import labelLine, labelLines
 
@@ -203,22 +201,22 @@ def test_nan_warning():
 
     line = plt.plot(x, y, label="test")[0]
 
-    with warnings.catch_warnings(record=True) as w:
+    warn_msg = (
+        ".* could not be annotated due to `nans` values. "
+        "Consider using another location via the `x` argument."
+    )
+    with pytest.warns(UserWarning, match=warn_msg):
         labelLine(line, 0.5)
-        assert issubclass(w[-1].category, UserWarning)
-        assert "could not be annotated" in str(w[-1].message)
 
-    with warnings.catch_warnings(record=True) as w:
-        labelLine(line, 2.5)
-        assert len(w) == 0
+    labelLine(line, 2.5)
 
 
 def test_nan_failure():
     x = np.array([0, 1])
     y = np.array([np.nan, np.nan])
 
     line = plt.plot(x, y, label="test")[0]
-    with assert_raises(Exception):
+    with pytest.raises(Exception):
         labelLine(line, 0.5)
 
 
@@ -228,9 +226,9 @@ def test_label_range(setupMpl):
     line = plt.plot(x, x**2, label="lorem ipsum")[0]
 
     # This should fail
-    with assert_raises(Exception):
+    with pytest.raises(Exception):
         labelLine(line, -1)
-    with assert_raises(Exception):
+    with pytest.raises(Exception):
         labelLine(line, 2)
 
     # This should work
@@ -365,3 +363,44 @@ def test_errorbars(setupMpl):
 
     labelLines(ax.get_lines(), align=False, xvals=pos)
     return fig
+
+
+@pytest.fixture
+def create_plot():
+    fig, ax = plt.subplots()
+    X = [0, 1]
+    Y = [0, 1]
+
+    lines = (
+        *ax.plot(X, Y, label="label1"),
+        *ax.plot(X, Y),  # no label
+        *ax.plot(X, Y, label="label2"),
+    )
+    return fig, ax, lines
+
+
+def test_warning_line_labeling(create_plot):
+    _fig, _ax, lines = create_plot
+
+    warn_msg = "Tried to label line .*, but could not find a label for it."
+    with pytest.warns(UserWarning, match=warn_msg):
+        txts = labelLines(lines)
+    # Make sure only two lines have been labeled
+    assert len(txts) == 2
+
+    with pytest.warns(UserWarning, match=warn_msg):
+        txts = labelLines(lines[1:])
+    # Make sure only one line has been labeled
+    assert len(txts) == 1
+
+
+def test_no_warning_line_labeling(create_plot):
+    _fig, _ax, lines = create_plot
+
+    txts = labelLines(lines[0:1])
+    assert len(txts) == 1
+
+
+def test_labeling_by_axis(create_plot):
+    txts = labelLines()
+    assert len(txts) == 2

pytest.ini

@@ -1,2 +1,3 @@
 [pytest]
 python_files = test*.py
+filterwarnings = error