Add test for datetime and categorical

l-johnston · l-johnston · commit fb76ed7464fc · 2020-08-21T14:01:48.000-05:00
diff --git a/lib/matplotlib/tests/test_units.py b/lib/matplotlib/tests/test_units.py
@@ -1,10 +1,11 @@
-from datetime import datetime
+from datetime import datetime, timezone, timedelta
 import platform
 from unittest.mock import MagicMock
 
 import matplotlib.pyplot as plt
 from matplotlib.testing.decorators import check_figures_equal, image_comparison
 import matplotlib.units as munits
+from matplotlib.category import UnitData
 import numpy as np
 import pytest
 
@@ -188,3 +189,17 @@ def test_set_xyunits(quantity_converter):
     ax2.yaxis.set_units("inches")
     assert ax1.xaxis.get_units() == ax2.xaxis.get_units() == "seconds"
     assert ax1.yaxis.get_units() == ax2.yaxis.get_units() == "inches"
+    y1 = [datetime(2020, i, 1, tzinfo=timezone.utc) for i in range(1, 13)]
+    y2 = [datetime(2021, i, 1, tzinfo=timezone.utc) for i in range(1, 13)]
+    fig, (ax1, ax2) = plt.subplots(1, 2, sharey=True)
+    ax1.plot(y1)
+    ax2.plot(y2)
+    ax1.yaxis.set_units(timezone(timedelta(hours=5)))
+    assert ax2.yaxis.units == timezone(timedelta(hours=5))
+    d1 = {"a": 1, "b": 2}
+    d2 = {"a": 3, "b": 4}
+    fig, (ax1, ax2) = plt.subplots(1, 2, sharex=True, sharey=True)
+    ax1.plot(d1.keys(), d1.values())
+    ax2.plot(d2.keys(), d2.values())
+    ax1.xaxis.set_units(UnitData(["c", "d"]))
+    assert "c" in ax2.xaxis.get_units()._mapping.keys()
