Treat units info as nu, make nu a list rather than single function

Author: ksunden

data_prototype/wrappers.py

@@ -1,4 +1,5 @@
 from typing import List, Dict, Any, Protocol, Tuple, get_type_hints
+import inspect
 
 import numpy as np
 
@@ -96,7 +97,7 @@ def draw(self, renderer):
     def _update_wrapped(self, data):
         raise NotImplementedError
 
-    def _query_and_transform(self, renderer, *, xunits: List[str], yunits: List[str]) -> Dict[str, Any]:
+    def _query_and_transform(self, renderer) -> Dict[str, Any]:
         """
         Helper to centralize the data querying and python-side transforms
 
@@ -126,27 +127,39 @@ def _query_and_transform(self, renderer, *, xunits: List[str], yunits: List[str]
             return self._cache[cache_key]
         except KeyError:
             ...
-        # TODO decide if units go pre-nu or post-nu?
-        for x_like in xunits:
-            if x_like in data:
-                data[x_like] = ax.xaxis.convert_units(data[x_like])
-        for y_like in yunits:
-            if y_like in data:
-                data[y_like] = ax.xaxis.convert_units(data[y_like])
-
         # doing the nu work here is nice because we can write it once, but we
         # really want to push this computation down a layer
         # TODO sort out how this interoperates with the transform stack
-        data = {k: self.nus.get(k, lambda x: x)(v) for k, v in data.items()}
+        def linearize(funclist):
+            if inspect.isfunction(funclist):
+                return funclist
+            def ret(x):
+                for func in funclist:
+                    x = func(x)
+                return x
+            return ret
+        data = {k: linearize(self.nus.get(k, lambda x: x))(v) for k, v in data.items()}
         self._cache[cache_key] = data
         return data
 
-    def __init__(self, data, nus, **kwargs):
+    def __init__(self, data, nus, xunits: List[str] = [], yunits: List[str] = [], **kwargs):
         super().__init__(**kwargs)
         self.data = data
         self._cache = LFUCache(64)
         # TODO make sure mutating this will invalidate the cache!
         self.nus = nus or {}
+        for field in xunits:
+            if field not in self.nus:
+                self.nus[field] = []
+            if inspect.isfunction(self.nus[field]):
+                self.nus[field] = [self.nus[field]]
+            self.nus[field].append(lambda x: self.axes.xaxis.convert_units(x))
+        for field in yunits:
+            if field not in self.nus:
+                self.nus[field] = []
+            if inspect.isfunction(self.nus[field]):
+                self.nus[field] = [self.nus[field]]
+            self.nus[field].append(lambda y: self.axes.yaxis.convert_units(y))
         self.stale = True
 
 
@@ -176,13 +189,13 @@ class LineWrapper(ProxyWrapper):
     _privtized_methods = ("set_xdata", "set_ydata", "set_data", "get_xdata", "get_ydata", "get_data")
 
     def __init__(self, data: DataContainer, nus=None, /, **kwargs):
-        super().__init__(data, nus)
+        super().__init__(data, nus, xunits=["x"], yunits=["y"])
         self._wrapped_instance = self._wrapped_class([], [], **kwargs)
 
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(renderer, xunits=["x"], yunits=["y"]),
+            self._query_and_transform(renderer),
         )
         return self._wrapped_instance.draw(renderer)
 
@@ -204,14 +217,14 @@ def __init__(self, data: DataContainer, nus=None, /, cmap=None, norm=None, **kwa
             if norm is None:
                 raise ValueError("not sure how to do autoscaling yet")
             nus["image"] = lambda image: cmap(norm(image))
-        super().__init__(data, nus)
+        super().__init__(data, nus, xunits=["xextent"], yunits=["yextent"])
         kwargs.setdefault("origin", "lower")
         self._wrapped_instance = self._wrapped_class(None, **kwargs)
 
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(renderer, xunits=["xextent"], yunits=["yextent"]),
+            self._query_and_transform(renderer),
         )
         return self._wrapped_instance.draw(renderer)
 
@@ -225,13 +238,13 @@ class StepWrapper(ProxyWrapper):
     _privtized_methods = ()  # ("set_data", "get_data")
 
     def __init__(self, data: DataContainer, nus=None, /, **kwargs):
-        super().__init__(data, nus)
+        super().__init__(data, nus, xunits=["edges"], yunits=["density"])
         self._wrapped_instance = self._wrapped_class([], [1], **kwargs)
 
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(renderer, xunits=["edges"], yunits=["density"]),
+            self._query_and_transform(renderer),
         )
         return self._wrapped_instance.draw(renderer)
 
@@ -251,7 +264,7 @@ def __init__(self, data: DataContainer, format_func, nus=None, /, **kwargs):
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(renderer, xunits=[], yunits=[]),
+            self._query_and_transform(renderer),
         )
         return self._wrapped_instance.draw(renderer)
 
@@ -297,7 +310,7 @@ def get_children(self):
 
 class ErrorbarWrapper(MultiProxyWrapper):
     def __init__(self, data: DataContainer, nus=None, /, **kwargs):
-        super().__init__(data, nus)
+        super().__init__(data, nus, xunits=["x", "xupper", "xlower"], yunits=["y", "yupper", "ylower"])
         # TODO all of the kwarg teasing apart that is needed
         color = kwargs.pop("color", "k")
         lw = kwargs.pop("lw", 2)
@@ -325,7 +338,7 @@ def __init__(self, data: DataContainer, nus=None, /, **kwargs):
     def draw(self, renderer):
         self._update_wrapped(
             self._query_and_transform(
-                renderer, xunits=["x", "xupper", "xlower"], yunits=["y", "yupper", "ylower"]
+                renderer
             ),
         )
         for k, v in self._wrapped_instances.items():