Implement DelayedConversionNode as a generalization on evaluate-time dependency injection

Used for units initially, presented as an alternative to #34 which provides more general tools

Author: ksunden

data_prototype/axes.py

@@ -10,7 +10,7 @@
 
 from .containers import ArrayContainer, DataUnion
 from .conversion_node import (
-    MatplotlibUnitConversion,
+    DelayedConversionNode,
     FunctionConversionNode,
     RenameConversionNode,
 )
@@ -96,8 +96,8 @@ def scatter(
         cont = DataUnion(defaults, inputs)
 
         pipeline = []
-        xconvert = MatplotlibUnitConversion.from_keys(("x",), axis=self.xaxis)
-        yconvert = MatplotlibUnitConversion.from_keys(("y",), axis=self.yaxis)
+        xconvert = DelayedConversionNode.from_keys(("x",), converter_key="xunits")
+        yconvert = DelayedConversionNode.from_keys(("y",), converter_key="yunits")
         pipeline.extend([xconvert, yconvert])
         pipeline.append(lambda x: np.ma.ravel(x))
         pipeline.append(lambda y: np.ma.ravel(y))

data_prototype/conversion_node.py

@@ -6,18 +6,22 @@
 import inspect
 from functools import cached_property
 
-from matplotlib.axis import Axis
-
 from typing import Any
 
 
-def evaluate_pipeline(nodes: Sequence[ConversionNode], input: dict[str, Any]):
+def evaluate_pipeline(
+    nodes: Sequence[ConversionNode],
+    input: dict[str, Any],
+    delayed_converters: dict[str, Callable] | None = None,
+):
     for node in nodes:
         if isinstance(node, Callable):
             k = list(inspect.signature(node).parameters.keys())[0]
             node = FunctionConversionNode.from_funcs({k: node})
-
-        input = node.evaluate(input)
+        if isinstance(node, DelayedConversionNode):
+            input = node.evaluate(input, delayed_converters)
+        else:
+            input = node.evaluate(input)
     return input
 
 
@@ -122,17 +126,24 @@ def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
 
 
 @dataclass
-class MatplotlibUnitConversion(ConversionNode):
-    axis: Axis
+class DelayedConversionNode(ConversionNode):
+    converter_key: str
 
     @classmethod
-    def from_keys(cls, keys: Sequence[str], axis: Axis):
-        return cls(tuple(keys), tuple(keys), trim_keys=False, axis=axis)
+    def from_keys(cls, keys: Sequence[str], converter_key: str):
+        return cls(
+            tuple(keys), tuple(keys), trim_keys=False, converter_key=converter_key
+        )
 
-    def evaluate(self, input: dict[str, Any]) -> dict[str, Any]:
+    def evaluate(
+        self, input: dict[str, Any], converters: dict[str, Callable] | None = None
+    ) -> dict[str, Any]:
         return super().evaluate(
             {
                 **input,
-                **{k: self.axis.convert_units(input[k]) for k in self.required_keys},
+                **{
+                    k: converters[self.converter_key](input[k])
+                    for k in self.required_keys
+                },
             }
         )

data_prototype/wrappers.py

@@ -1,4 +1,4 @@
-from typing import List, Dict, Any, Protocol, Tuple, get_type_hints
+from typing import Dict, Any, Protocol, Tuple, get_type_hints
 import inspect
 
 import numpy as np
@@ -121,17 +121,13 @@ 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
 
         Parameters
         ----------
         renderer : RendererBase
-        xunits, yunits : List[str]
-            The list of keys that need to be run through the x and y unit machinery.
         """
         # extract what we need to about the axes to query the data
         ax = self.axes
@@ -153,8 +149,11 @@ def _query_and_transform(
             return self._cache[cache_key]
         except KeyError:
             ...
-        # TODO units
-        transformed_data = evaluate_pipeline(self._converters, data)
+        delayed_conversion = {
+            "xunits": ax.xaxis.convert_units,
+            "yunits": ax.yaxis.convert_units,
+        }
+        transformed_data = evaluate_pipeline(self._converters, data, delayed_conversion)
 
         self._cache[cache_key] = transformed_data
         return transformed_data
@@ -232,7 +231,7 @@ def __init__(self, data: DataContainer, converters=None, /, **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)
 
@@ -265,7 +264,7 @@ def __init__(self, data: DataContainer, converters=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)
 
@@ -304,7 +303,7 @@ def __init__(
     @_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)
 
@@ -325,7 +324,7 @@ def __init__(self, data: DataContainer, converters=None, /, **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)
 
@@ -344,7 +343,7 @@ def __init__(self, data: DataContainer, converters=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)
 
@@ -425,11 +424,7 @@ def __init__(self, data: DataContainer, converters=None, /, **kwargs):
     @_stale_wrapper
     def draw(self, renderer):
         self._update_wrapped(
-            self._query_and_transform(
-                renderer,
-                xunits=["x", "xupper", "xlower"],
-                yunits=["y", "yupper", "ylower"],
-            ),
+            self._query_and_transform(renderer),
         )
         for k, v in self._wrapped_instances.items():
             v.draw(renderer)

examples/units.py

@@ -12,7 +12,7 @@
 import matplotlib.markers as mmarkers
 
 from data_prototype.containers import ArrayContainer
-from data_prototype.conversion_node import MatplotlibUnitConversion
+from data_prototype.conversion_node import DelayedConversionNode
 
 from data_prototype.wrappers import PathCollectionWrapper
 
@@ -35,8 +35,10 @@
 fig, ax = plt.subplots()
 ax.set_xlim(-0.5, 7)
 ax.set_ylim(0, 5)
-conv = MatplotlibUnitConversion.from_keys(("x",), axis=ax.xaxis)
-lw = PathCollectionWrapper(cont, [conv], offset_transform=ax.transData)
+xconv = DelayedConversionNode.from_keys(("x",), converter_key="xunits")
+yconv = DelayedConversionNode.from_keys(("y",), converter_key="yunits")
+lw = PathCollectionWrapper(cont, [xconv, yconv], offset_transform=ax.transData)
 ax.add_artist(lw)
 ax.xaxis.set_units(ureg.feet)
+ax.yaxis.set_units(ureg.m)
 plt.show()