restored test_from_category_tests

story645 · story645 · commit 5ff0dbfd68b2 · 2016-08-17T15:17:13.000-04:00
diff --git a/lib/matplotlib/category.py b/lib/matplotlib/category.py
@@ -113,35 +113,28 @@ class CategoryNorm(mcolors.Normalize):
     """
     Preserves ordering of discrete values
     """
-    def __init__(self, categories):
+    def __init__(self, data):
         """
         *categories*
             distinct values for mapping
 
         Out-of-range values are mapped to np.nan
         """
 
-        # facilitates cleaner DuckTyping of axis interface
+        self.units = StrCategoryConverter()
+        self.unit_data = None
+        self.units.default_units(data,
+                                 self, sort=False)
 
-        class CatAxis(object):
-            def __init__(self):
-                self.unit_data = None
-                self.units = StrCategoryConverter()
-
-        self.axis = CatAxis()
-        self.axis.units.default_units(categories, self.axis,
-                                      sort=False)
-
-        nvals = self.axis.unit_data.locs
-        self.vmin = min(nvals)
-        self.vmax = max(nvals)
+        self.vmin = min(self.unit_data.locs)
+        self.vmax = max(self.unit_data.locs)
 
     def __call__(self, value, clip=None):
         # gonna have to go into imshow and undo casting
         value = np.asarray(value, dtype=np.int)
-        ret = self.axis.units.convert(value, None, self.axis)
+        ret = self.units.convert(value, None, self)
         # knock out values not in the norm
-        mask = np.in1d(ret, self.axis.unit_data.locs).reshape(ret.shape)
+        mask = np.in1d(ret, self.unit_data.locs).reshape(ret.shape)
         # normalize ret & locs
         ret /= self.vmax
         return np.ma.array(ret, mask=~mask)
@@ -162,8 +155,7 @@ def colors_from_categories(codings):
                    :class:`Normalize` instance
     """
     if isinstance(codings, dict):
-        codings = codings.items()
-
+        codings = cbook.sanitize_sequence(codings.items())
     values, colors = zip(*codings)
     cmap = mcolors.ListedColormap(list(colors))
     norm = CategoryNorm(list(values))
diff --git a/lib/matplotlib/tests/test_category.py b/lib/matplotlib/tests/test_category.py
@@ -141,11 +141,15 @@ def test_norm(self, data, nmap):
         np.testing.assert_allclose(norm(data), nmap)
 
 
-def test_colors_from_categories():
-    codings = {205: "red", 101: "blue", 302: "green"}
-    cmap, norm = cat.colors_from_categories(codings)
-    assert cmap.colors == ['red', 'green', 'blue']
-    assert norm.axis.unit_data.seq == ['205', '302', '101']
+class TestColorsFromCategories(object):
+    testdata = [[{'101': "blue", '205': "red", '302': "green"}, dict],
+                [[('205', "red"), ('101', "blue"), ('302', "green")], list]]
+    ids = ["dict", "tuple"]
+
+    @pytest.mark.parametrize("codings, mtype", testdata, ids=ids)
+    def test_colors_from_categories(self, codings, mtype):
+        cmap, norm = cat.colors_from_categories(codings)
+        assert mtype(zip(norm.unit_data.seq, cmap.colors)) == codings
 
 
 def lt(tl):
