BUG: modified safe_first_element and added tests

渡邉 美希 · tacaswell · 渡邉 美希 · commit a8c01a42c5bb · 2022-08-12T03:34:10.000+09:00
Co-authored-by: Thomas A Caswell &lt;tcaswell@gmail.com&gt;
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -2176,12 +2176,12 @@ def _convert_dx(dx, x0, xconv, convert):
             # removes the units from unit packages like `pint` that
             # wrap numpy arrays.
             try:
-                x0 = cbook.safe_first_element(x0)
+                x0 = cbook._safe_first_non_none(x0)
             except (TypeError, IndexError, KeyError):
                 pass
 
             try:
-                x = cbook.safe_first_element(xconv)
+                x = cbook._safe_first_non_none(xconv)
             except (TypeError, IndexError, KeyError):
                 x = xconv
 
@@ -2779,11 +2779,11 @@ def broken_barh(self, xranges, yrange, **kwargs):
         """
         # process the unit information
         if len(xranges):
-            xdata = cbook.safe_first_element(xranges)
+            xdata = cbook._safe_first_non_none(xranges)
         else:
             xdata = None
         if len(yrange):
-            ydata = cbook.safe_first_element(yrange)
+            ydata = cbook._safe_first_non_none(yrange)
         else:
             ydata = None
         self._process_unit_info(
@@ -3424,10 +3424,10 @@ def _upcast_err(err):
                     # safe_first_element because getitem is index-first not
                     # location first on pandas objects so err[0] almost always
                     # fails.
-                    isinstance(cbook.safe_first_element(err), np.ndarray)
+                    isinstance(cbook._safe_first_non_none(err), np.ndarray)
             ):
                 # Get the type of the first element
-                atype = type(cbook.safe_first_element(err))
+                atype = type(cbook._safe_first_non_none(err))
                 # Promote the outer container to match the inner container
                 if atype is np.ndarray:
                     # Converts using np.asarray, because data cannot
@@ -4290,7 +4290,7 @@ def _parse_scatter_color_args(c, edgecolors, kwargs, xsize,
         c_is_string_or_strings = (
             isinstance(c, str)
             or (np.iterable(c) and len(c) > 0
-                and isinstance(cbook.safe_first_element(c), str)))
+                and isinstance(cbook._safe_first_non_none(c), str)))
 
         def invalid_shape_exception(csize, xsize):
             return ValueError(
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -1688,22 +1688,42 @@ def safe_first_element(obj):
     """
     Return the first element in *obj*.
 
-    This is an type-independent way of obtaining the first element, supporting
-    both index access and the iterator protocol.
-    """
-    if isinstance(obj, collections.abc.Iterator):
-        # needed to accept `array.flat` as input.
-        # np.flatiter reports as an instance of collections.Iterator
-        # but can still be indexed via [].
-        # This has the side effect of re-setting the iterator, but
-        # that is acceptable.
-        try:
-            return obj[0]
-        except TypeError:
-            pass
-        raise RuntimeError("matplotlib does not support generators "
-                           "as input")
-    return next(iter(obj))
+    This is an type-independent way of obtaining the first element,
+    supporting both index access and the iterator protocol.
+    """
+    return _safe_first_non_none(obj, skip_none=False)
+
+
+def _safe_first_non_none(obj, skip_none=True):
+    """
+    Return the first non-None element in *obj*.
+    This is a method for internal use.
+
+    This is an type-independent way of obtaining the first non-None element,
+    supporting both index access and the iterator protocol.
+    The first non-None element will be obtained when skip_none is True.
+    """
+    if skip_none is False:
+        if isinstance(obj, collections.abc.Iterator):
+            # needed to accept `array.flat` as input.
+            # np.flatiter reports as an instance of collections.Iterator
+            # but can still be indexed via [].
+            # This has the side effect of re-setting the iterator, but
+            # that is acceptable.
+            try:
+                return obj[0]
+            except TypeError:
+                pass
+            raise RuntimeError("matplotlib does not support generators "
+                               "as input")
+        return next(iter(obj))
+    elif isinstance(obj, np.flatiter):
+        return obj[0]
+    elif isinstance(obj, collections.abc.Iterator):
+        raise RuntimeError("matplotlib does not "
+                           "support generators as input")
+    else:
+        return next(val for val in obj if val is not None)
 
 
 def sanitize_sequence(data):
diff --git a/lib/matplotlib/dates.py b/lib/matplotlib/dates.py
@@ -1908,7 +1908,7 @@ def default_units(x, axis):
             x = x.ravel()
 
         try:
-            x = cbook.safe_first_element(x)
+            x = cbook._safe_first_non_none(x)
         except (TypeError, StopIteration):
             pass
 
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -4,7 +4,7 @@
 from weakref import ref
 from unittest.mock import patch, Mock
 
-from datetime import datetime
+from datetime import datetime, date, timedelta
 
 import numpy as np
 from numpy.testing import (assert_array_equal, assert_approx_equal,
@@ -602,7 +602,7 @@ def test_flatiter():
     it = x.flat
     assert 0 == next(it)
     assert 1 == next(it)
-    ret = cbook.safe_first_element(it)
+    ret = cbook._safe_first_non_none(it)
     assert ret == 0
 
     assert 0 == next(it)
@@ -758,7 +758,7 @@ def test_contiguous_regions():
 def test_safe_first_element_pandas_series(pd):
     # deliberately create a pandas series with index not starting from 0
     s = pd.Series(range(5), index=range(10, 15))
-    actual = cbook.safe_first_element(s)
+    actual = cbook._safe_first_non_none(s)
     assert actual == 0
 
 
@@ -888,3 +888,10 @@ def test_format_approx():
     assert f(0.0012345600001, 5) == '0.00123'
     assert f(-0.0012345600001, 5) == '-0.00123'
     assert f(0.0012345600001, 8) == f(0.0012345600001, 10) == '0.00123456'
+
+
+def test_safe_first_element_with_none():
+    datetime_lst = [date.today() + timedelta(days=i) for i in range(10)]
+    datetime_lst[0] = None
+    actual = cbook._safe_first_non_none(datetime_lst)
+    assert actual is not None and actual == datetime_lst[1]
diff --git a/lib/matplotlib/units.py b/lib/matplotlib/units.py
@@ -197,7 +197,7 @@ def get_converter(self, x):
             except KeyError:
                 pass
         try:  # If cache lookup fails, look up based on first element...
-            first = cbook.safe_first_element(x)
+            first = cbook._safe_first_non_none(x)
         except (TypeError, StopIteration):
             pass
         else:
