created dummy axis on norm to store conversion machinary

story645 · story645 · commit 748ff75a025e · 2016-08-17T11:15:22.000-04:00
diff --git a/lib/matplotlib/category.py b/lib/matplotlib/category.py
@@ -74,15 +74,14 @@ def axisinfo(unit, axis):
         return munits.AxisInfo(majloc=majloc, majfmt=majfmt)
 
     @staticmethod
-    def default_units(data, axis, sort=True):
+    def default_units(data, axis, sort=True, normed=False):
         """
         Create mapping between string categories in *data*
-        and integers, then store in *axis.unit_data*
+        and integers, and store in *axis.unit_data*
         """
-
         if axis and axis.unit_data:
             axis.unit_data.update(data, sort)
-            return
+            return axis.unit_data
 
         unit_data = UnitData(data, sort)
         if axis:
@@ -122,22 +121,29 @@ def __init__(self, categories):
         Out-of-range values are mapped to np.nan
         """
 
-        self.unit_data = StrCategoryConverter.default_units(categories,
-                                                            None, sort=False)
-        self.categories = to_str_array(categories)
-        self.N = len(self.categories)
-        self.nvals = self.unit_data.locs
-        self.vmin = min(self.nvals)
-        self.vmax = max(self.nvals)
+        # facilitates cleaner DuckTyping of axis interface
+
+        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)
 
     def __call__(self, value, clip=None):
         # gonna have to go into imshow and undo casting
-        value = np.asarray(value, dtype=int)
-        ret = StrCategoryConverter.convert(value, None, self)
+        value = np.asarray(value, dtype=np.int)
+        ret = self.axis.units.convert(value, None, self.axis)
         # knock out values not in the norm
-        mask = np.in1d(ret, self.unit_data.locs).reshape(ret.shape)
-        # normalize ret
-        ret /=  self.vmax
+        mask = np.in1d(ret, self.axis.unit_data.locs).reshape(ret.shape)
+        # normalize ret & locs
+        ret /= self.vmax
         return np.ma.array(ret, mask=~mask)
 
 
@@ -184,7 +190,7 @@ def convert_to_string(value):
 
 
 class UnitData(object):
-    # debatable makes sense to special code missing values
+    # debatable if it makes sense to special code missing values
     spdict = {'nan': -1.0, 'inf': -2.0, '-inf': -3.0}
 
     def __init__(self, data, sort=True):
@@ -202,7 +208,7 @@ def __init__(self, data, sort=True):
         self._set_seq_locs(data, 0, sort)
         self.sort = sort
 
-    def update(self, new_data, sort=None):
+    def update(self, new_data, sort=True):
         if sort:
             self.sort = sort
         # so as not to conflict with spdict
diff --git a/lib/matplotlib/tests/test_category.py b/lib/matplotlib/tests/test_category.py
@@ -129,24 +129,23 @@ def test_StrCategoryFormatterUnicode(self):
 
 
 class TestCategoryNorm(object):
-    testdata = [[[205, 302, 205, 101], [0, 2, 0, 1]],
-                [[205, np.nan, 101, 305], [0, np.nan, 1, 2]],
-                [[205, 101, 504, 101], [0, 1, np.nan, 1]]]
+    testdata = [[[205, 302, 205, 101], [0, 1, 0, .5]],
+                [[205, np.nan, 101, 305], [0, np.nan, .5, 1]],
+                [[205, 101, 504, 101], [0, .5, np.nan, .5]]]
 
     ids = ["regular", "nan", "exclude"]
 
     @pytest.mark.parametrize("data, nmap", testdata, ids=ids)
     def test_norm(self, data, nmap):
         norm = cat.CategoryNorm([205, 101, 302])
-        np.testing.assert_array_equal(norm(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']
-    np.testing.assert_array_equal(norm.categories, ['205', '302', '101'])
-    assert cmap.N == norm.N
+    assert norm.axis.unit_data.seq == ['205', '302', '101']
 
 
 def lt(tl):
