Merge pull request #9314 from anntzer/simpler-units

timhoffm · web-flow · commit cd3ed7acdad0 · 2019-02-28T17:48:01.000+01:00
Simplify units.Registry.get_converter.
diff --git a/lib/matplotlib/units.py b/lib/matplotlib/units.py
@@ -55,9 +55,10 @@ class ConversionError(TypeError):
 
 class AxisInfo(object):
     """
-    Information to support default axis labeling, tick labeling, and
-    default limits. An instance of this class must be returned by
-    :meth:`ConversionInterface.axisinfo`.
+    Information to support default axis labeling, tick labeling, and limits.
+
+    An instance of this class must be returned by
+    `ConversionInterface.axisinfo`.
     """
     def __init__(self, majloc=None, minloc=None,
                  majfmt=None, minfmt=None, label=None,
@@ -96,8 +97,7 @@ class ConversionInterface(object):
     @staticmethod
     def axisinfo(unit, axis):
         """
-        Return an `~units.AxisInfo` instance for the axis with the
-        specified units.
+        Return an `~units.AxisInfo` for the axis with the specified units.
         """
         return None
 
@@ -112,19 +112,19 @@ def default_units(x, axis):
     def convert(obj, unit, axis):
         """
         Convert *obj* using *unit* for the specified *axis*.
-        If *obj* is a sequence, return the converted sequence.
-        The output must be a sequence of scalars that can be used by the numpy
-        array layer.
+
+        If *obj* is a sequence, return the converted sequence.  The output must
+        be a sequence of scalars that can be used by the numpy array layer.
         """
         return obj
 
     @staticmethod
     def is_numlike(x):
         """
-        The Matplotlib datalim, autoscaling, locators etc work with
-        scalars which are the units converted to floats given the
-        current unit.  The converter may be passed these floats, or
-        arrays of them, even when units are set.
+        The Matplotlib datalim, autoscaling, locators etc work with scalars
+        which are the units converted to floats given the current unit.  The
+        converter may be passed these floats, or arrays of them, even when
+        units are set.
         """
         if np.iterable(x):
             for thisx in x:
@@ -134,73 +134,33 @@ def is_numlike(x):
 
 
 class Registry(dict):
-    """
-    A register that maps types to conversion interfaces.
-    """
-    def __init__(self):
-        dict.__init__(self)
-        self._cached = {}
+    """Register types with conversion interface."""
 
     def get_converter(self, x):
-        """
-        Get the converter for data that has the same type as *x*. If no
-        converters are registered for *x*, returns ``None``.
-        """
-
-        if not len(self):
-            return None  # nothing registered
-        # DISABLED idx = id(x)
-        # DISABLED cached = self._cached.get(idx)
-        # DISABLED if cached is not None: return cached
-
-        converter = None
-        classx = getattr(x, '__class__', None)
-
-        if classx is not None:
-            converter = self.get(classx)
-
-        if converter is None and hasattr(x, "values"):
-            # this unpacks pandas series or dataframes...
-            x = x.values
-
-        # If x is an array, look inside the array for data with units
+        """Get the converter interface instance for *x*, or None."""
+        if hasattr(x, "values"):
+            x = x.values  # Unpack pandas Series and DataFrames.
         if isinstance(x, np.ndarray):
+            # In case x in a masked array, access the underlying data (only its
+            # type matters).  If x is a regular ndarray, getdata() just returns
+            # the array itself.
+            x = np.ma.getdata(x).ravel()
             # If there are no elements in x, infer the units from its dtype
             if not x.size:
                 return self.get_converter(np.array([0], dtype=x.dtype))
-            xravel = x.ravel()
-            try:
-                # pass the first value of x that is not masked back to
-                # get_converter
-                if not np.all(xravel.mask):
-                    # Get first non-masked item
-                    converter = self.get_converter(
-                        xravel[np.argmin(xravel.mask)])
-                    return converter
-            except AttributeError:
-                # not a masked_array
-                # Make sure we don't recurse forever -- it's possible for
-                # ndarray subclasses to continue to return subclasses and
-                # not ever return a non-subclass for a single element.
-                next_item = xravel[0]
-                if (not isinstance(next_item, np.ndarray) or
-                        next_item.shape != x.shape):
-                    converter = self.get_converter(next_item)
-                return converter
-
-        # If we haven't found a converter yet, try to get the first element
-        if converter is None:
-            try:
-                thisx = cbook.safe_first_element(x)
+        try:  # Look up in the cache.
+            return self[type(x)]
+        except KeyError:
+            try:  # If cache lookup fails, look up based on first element...
+                first = cbook.safe_first_element(x)
             except (TypeError, StopIteration):
                 pass
             else:
-                if classx and classx != getattr(thisx, '__class__', None):
-                    converter = self.get_converter(thisx)
-                    return converter
-
-        # DISABLED self._cached[idx] = converter
-        return converter
+                # ... and avoid infinite recursion for pathological iterables
+                # where indexing returns instances of the same iterable class.
+                if type(first) is not type(x):
+                    return self.get_converter(first)
+        return None
 
 
 registry = Registry()
