API: restore support for bytes

tacaswell · tacaswell · commit cd31afd8b2be · 2018-02-04T19:54:41.000-05:00
diff --git a/lib/matplotlib/category.py b/lib/matplotlib/category.py
@@ -17,14 +17,16 @@ class StrCategoryConverter(units.ConversionInterface):
     @staticmethod
     def convert(value, unit, axis):
         """Use axis.units mapping to map categorical data to floats."""
-
+        def getter(k):
+            if not isinstance(k, six.text_type):
+                k = k.decode('utf-8')
+            return axis.units._mapping[k]
         # We also need to pass numbers through.
         if np.issubdtype(np.asarray(value).dtype.type, np.number):
             return value
         else:
             axis.units.update(value)
-            str2idx = np.vectorize(axis.units._mapping.__getitem__,
-                                   otypes=[float])
+            str2idx = np.vectorize(getter, otypes=[float])
             return str2idx(value)
 
     @staticmethod
@@ -61,6 +63,9 @@ def __call__(self, x, pos=None):
 
 
 class UnitData(object):
+    valid_types = tuple(set(six.string_types +
+                            (bytes, six.text_type, np.str_, np.bytes_)))
+
     def __init__(self, data=None):
         """Create mapping between unique categorical values and numerical id.
 
@@ -73,10 +78,12 @@ def __init__(self, data=None):
         self._vals = []
         if data is None:
             data = ()
-        self._mapping = OrderedDict(data)
-        for k, v in self._mapping.items():
-            if not isinstance(k, six.text_type):
+        self._mapping = OrderedDict()
+        for k, v in OrderedDict(data).items():
+            if not isinstance(k, self.valid_types):
                 raise TypeError("{val!r} is not a string".format(val=k))
+            if not isinstance(k, six.text_type):
+                k = k.decode('utf-8')
             self._mapping[k] = int(v)
         if self._mapping:
             start = max(self._mapping.values()) + 1
@@ -85,19 +92,22 @@ def __init__(self, data=None):
         self._counter = itertools.count(start=start)
 
     def update(self, data):
-        if isinstance(data, six.string_types):
+        if isinstance(data, self.valid_types):
             data = [data]
         sorted_unique = OrderedDict.fromkeys(data)
         for val in sorted_unique:
+            if not isinstance(val, self.valid_types):
+                raise TypeError("{val!r} is not a string".format(val=val))
+            if not isinstance(val, six.text_type):
+                val = val.decode('utf-8')
             if val in self._mapping:
                 continue
-            if not isinstance(val, six.text_type):
-                raise TypeError("{val!r} is not a string".format(val=val))
             self._vals.append(val)
             self._mapping[val] = next(self._counter)
 
 
 # Connects the convertor to matplotlib
+
 units.registry[str] = StrCategoryConverter()
 units.registry[bytes] = StrCategoryConverter()
 units.registry[np.str_] = StrCategoryConverter()
diff --git a/lib/matplotlib/tests/test_category.py b/lib/matplotlib/tests/test_category.py
@@ -61,6 +61,8 @@ class TestStrCategoryConverter(object):
     test_cases = [("unicode", {u"Здравствуйте мир": 42}),
                   ("ascii", {"hello world": 42}),
                   ("single", {'a': 0, 'b': 1, 'c': 2}),
+                  ("single bytes", {b'a': 0, b'b': 1, b'c': 2}),
+                  ("mixed bytes", {b'a': 0, 'b': 1, b'c': 2}),
                   ("single + values>10", {'A': 0, 'B': 1, 'C': 2,
                                           'D': 3, 'E': 4, 'F': 5,
                                           'G': 6, 'H': 7, 'I': 8,
@@ -111,11 +113,11 @@ def test_StrCategoryFormatter(self):
         assert labels(1, 1) == "world"
 
     def test_StrCategoryFormatterUnicode(self):
-        seq = ["Здравствуйте", "привет"]
+        seq = [u"Здравствуйте", u"привет"]
         u = cat.UnitData()
         u.update(seq)
         labels = cat.StrCategoryFormatter(u)
-        assert labels(1, 1) == "привет"
+        assert labels(1, 1) == u"привет"
 
 
 def lt(tl):
@@ -130,6 +132,8 @@ def axis_test(axis, ticks, labels, unit_data):
 
 class TestBarsBytes(object):
     bytes_cases = [('string list', ['a', 'b', 'c']),
+                   ('bytes list', [b'a', b'b', b'c']),
+                   ('mixed list', [b'a', 'b', b'c']),
                    ]
 
     bytes_ids, bytes_data = zip(*bytes_cases)
@@ -232,6 +236,7 @@ def ax():
     [([u"Здравствуйте мир"], [0], [u"Здравствуйте мир"]),
      (["a", "b", "b", "a", "c", "c"], [0, 1, 1, 0, 2, 2], ["a", "b", "c"]),
      (["foo", "bar"], range(2), ["foo", "bar"]),
+     ([b"foo", "bar"], range(2), ["foo", "bar"]),
      (np.array(["1", "11", "3"]), range(3), ["1", "11", "3"])])
 def test_simple(ax, data, expected_indices, expected_labels):
     l, = ax.plot(data)
