Merge pull request #7241 from QuLogic/categorical-bytes

story645 · web-flow · commit 7fd4b69635ec · 2016-10-14T14:20:36.000-04:00
Categorical support for NumPy string arrays, maybe closes #7215
diff --git a/lib/matplotlib/category.py b/lib/matplotlib/category.py
@@ -131,7 +131,10 @@ def _set_seq_locs(self, data, value):
                 self.locs.append(value)
                 value += 1
 
+
 # Connects the convertor to matplotlib
 units.registry[str] = StrCategoryConverter()
-units.registry[bytes] = StrCategoryConverter()
+units.registry[np.str_] = StrCategoryConverter()
 units.registry[six.text_type] = StrCategoryConverter()
+units.registry[bytes] = StrCategoryConverter()
+units.registry[np.bytes_] = StrCategoryConverter()
diff --git a/lib/matplotlib/tests/test_category.py b/lib/matplotlib/tests/test_category.py
@@ -186,6 +186,41 @@ def test_plot_1d_missing(self):
 
         self.axis_test(ax.yaxis, self.dmticks, self.dmlabels, self.dmunit_data)
 
+    @cleanup
+    @pytest.mark.usefixtures("data")
+    @pytest.mark.parametrize("bars",
+                             [['a', 'b', 'c'],
+                              [b'a', b'b', b'c'],
+                              np.array([b'a', b'b', b'c'])],
+                             ids=['string list', 'bytes list',
+                                  'bytes ndarray'])
+    def test_plot_bytes(self, bars):
+        counts = np.array([4, 6, 5])
+
+        fig, ax = plt.subplots()
+        ax.bar(bars, counts)
+        fig.canvas.draw()
+
+        self.axis_test(ax.xaxis, self.dticks, self.dlabels, self.dunit_data)
+
+    @cleanup
+    @pytest.mark.parametrize("bars",
+                             [['1', '11', '3'],
+                              np.array(['1', '11', '3']),
+                              [b'1', b'11', b'3'],
+                              np.array([b'1', b'11', b'3'])],
+                             ids=['string list', 'string ndarray',
+                                  'bytes list', 'bytes ndarray'])
+    def test_plot_numlike(self, bars):
+        counts = np.array([4, 6, 5])
+
+        fig, ax = plt.subplots()
+        ax.bar(bars, counts)
+        fig.canvas.draw()
+
+        unitmap = MockUnitData([('1', 0), ('11', 1), ('3', 2)])
+        self.axis_test(ax.xaxis, [0, 1, 2], ['1', '11', '3'], unitmap)
+
     @cleanup
     @pytest.mark.usefixtures("data", "missing_data")
     def test_plot_2d(self):
