Make unit converters also handle instances of subclasses.

anntzer · anntzer · commit e756add5a4c7 · 2019-02-27T19:19:35.000+01:00
diff --git a/doc/users/next_whats_new/2019-02-27-AL.rst b/doc/users/next_whats_new/2019-02-27-AL.rst
@@ -0,0 +1,5 @@
+Unit converters now handle instances of subclasses
+``````````````````````````````````````````````````
+
+Unit converters now also handle instances of subclasses of the class they have
+been registered for.
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -1,10 +1,11 @@
+from datetime import datetime
+import platform
 from unittest.mock import MagicMock
 
 import matplotlib.pyplot as plt
-from matplotlib.testing.decorators import image_comparison
+from matplotlib.testing.decorators import check_figures_equal, image_comparison
 import matplotlib.units as munits
 import numpy as np
-import platform
 import pytest
 
 
@@ -121,7 +122,6 @@ def test_empty_set_limits_with_units(quantity_converter):
 @image_comparison(baseline_images=['jpl_bar_units'], extensions=['png'],
                   savefig_kwarg={'dpi': 120}, style='mpl20')
 def test_jpl_bar_units():
-    from datetime import datetime
     import matplotlib.testing.jpl_units as units
     units.register()
 
@@ -138,7 +138,6 @@ def test_jpl_bar_units():
 @image_comparison(baseline_images=['jpl_barh_units'], extensions=['png'],
                   savefig_kwarg={'dpi': 120}, style='mpl20')
 def test_jpl_barh_units():
-    from datetime import datetime
     import matplotlib.testing.jpl_units as units
     units.register()
 
@@ -155,3 +154,12 @@ def test_jpl_barh_units():
 def test_empty_arrays():
     # Check that plotting an empty array with a dtype works
     plt.scatter(np.array([], dtype='datetime64[ns]'), np.array([]))
+
+
+@check_figures_equal(extensions=["png"])
+def test_subclass(fig_test, fig_ref):
+    class subdate(datetime):
+        pass
+
+    fig_test.subplots().plot(subdate(2000, 1, 1), 0, "o")
+    fig_ref.subplots().plot(datetime(2000, 1, 1), 0, "o")
diff --git a/lib/matplotlib/units.py b/lib/matplotlib/units.py
@@ -149,8 +149,8 @@ def get_converter(self, x):
             if not x.size:
                 return self.get_converter(np.array([0], dtype=x.dtype))
         try:
-            return self[type(x)]
-        except KeyError:
+            return next(self[cls] for cls in type(x).__mro__ if cls in self)
+        except StopIteration:
             try:
                 first = cbook.safe_first_element(x)
             except (TypeError, StopIteration):
