Add axis labels feature to geonodes

Author: Griperis

data_vis/blends/chart_geonodes.blend

@@ -13,8 +13,6 @@
 from .data import DV_DataProperties
 from .components import (
     DV_AddAxis,
-    DV_AddHeading,
-    DV_AddAxisLabel,
     DV_AddDataLabels,
 )
 from .animations import (
@@ -35,8 +33,6 @@
     DV_GN_SurfaceChart,
     DV_GN_PieChart,
     DV_AddAxis,
-    DV_AddHeading,
-    DV_AddAxisLabel,
     DV_AddDataLabels,
     DV_ChartPanel,
     DV_AxisPanel,

data_vis/geonodes/__init__.py

@@ -156,6 +156,8 @@ def execute(self, context: bpy.types.Context):
         if self.axis_type == AxisType.CATEGORICAL:
             self._setup_categorical_axis(obj, mod)
 
+        self._add_labels_if_available(obj, mod)
+
         modifier_utils.add_used_materials_to_object(mod, obj)
         return {"FINISHED"}
 
@@ -180,12 +182,35 @@ def _setup_categorical_axis(
         assert is_chart(obj)
         data_from_obj = data.get_chart_data_info(obj)
         if data_from_obj is None:
-            logger.error("No data found on the chart {obj.name}")
+            logger.error(f"No data found on the chart {obj.name}")
             return
 
         modifier_utils.set_input(mod, "Tick Count", len(data_from_obj["categories"]))
         modifier_utils.set_input(mod, "Labels", ",".join(data_from_obj["categories"]))
 
+    def _add_labels_if_available(
+        self, obj: bpy.types.Object, mod: bpy.types.NodesModifier
+    ) -> None:
+        assert is_chart(obj)
+        data_from_obj = data.get_chart_data_info(obj)
+        if data_from_obj is None:
+            logger.error(f"No data found on the chart {obj.name}")
+            return
+
+        axis_labels = data_from_obj["axis_labels"]
+        if len(axis_labels) == 0:
+            return
+        if self.axis == "X" and len(axis_labels) > 0:
+            modifier_utils.set_input(mod, "Axis Label Text", axis_labels[0])
+
+        if self.axis == "Y" and len(axis_labels) > 1:
+            modifier_utils.set_input(mod, "Axis Label Text", axis_labels[1])
+
+        if self.axis == "Z" and len(axis_labels) == 2:
+            modifier_utils.set_input(mod, "Axis Label Text", axis_labels[1])
+        elif len(axis_labels) == 3:
+            modifier_utils.set_input(mod, "Axis Label Text", axis_labels[2])
+
 
 @data_vis_logging.logged_operator
 class DV_AddDataLabels(bpy.types.Operator):
@@ -204,27 +229,3 @@ def execute(self, context: bpy.types.Context):
         mod.node_group = library.load_above_data_labels()
         modifier_utils.add_used_materials_to_object(mod, obj)
         return {"FINISHED"}
-
-
-# TODO: Create a object that's in the middle of the axis and parented to the chart object
-@data_vis_logging.logged_operator
-class DV_AddAxisLabel(bpy.types.Operator):
-    bl_idname = "data_vis.add_axis_label"
-    bl_label = "Add Axis Label"
-    bl_options = {"REGISTER", "UNDO"}
-
-    # Adds a axis label to the selected chart
-    def execute(self, context):
-        return {"FINISHED"}
-
-
-# TODO: Add heading to the chart
-@data_vis_logging.logged_operator
-class DV_AddHeading(bpy.types.Operator):
-    # Adds a heading to the selected chart
-    bl_idname = "data_vis.add_heading"
-    bl_label = "Add Heading"
-    bl_options = {"REGISTER", "UNDO"}
-
-    def execute(self, context):
-        return {"FINISHED"}

data_vis/geonodes/components.py

@@ -135,6 +135,7 @@ class PreprocessedData:
     ws: np.ndarray | None = None
     z_ns: np.ndarray | None = None
     categories: np.ndarray | None = None
+    axis_labels: typing.List[str] = dataclasses.field(default_factory=list)
 
 
 def _store_chart_data_info(
@@ -144,8 +145,11 @@ def _store_chart_data_info(
         "data_type": data_type,
         "shape": verts.shape,
     }
-    if data is not None and data.categories is not None:
-        data_dict["categories"] = data.categories.tolist()
+    if data is not None:
+        if data.categories is not None:
+            data_dict["categories"] = data.categories.tolist()
+        if data.axis_labels is not None:
+            data_dict["axis_labels"] = data.axis_labels
 
     obj[DATA_TYPE_PROPERTY] = json.dumps(data_dict)
 
@@ -205,7 +209,7 @@ def _preprocess_data(data, data_type: str) -> PreprocessedData:
     else:
         raise RuntimeError(f"Unknown DataType {data_type}")
 
-    return PreprocessedData(vert_positions, ws, z_ns, categories)
+    return PreprocessedData(vert_positions, ws, z_ns, categories, axis_labels=[])
 
 
 @dataclasses.dataclass
@@ -220,7 +224,9 @@ def _convert_data_to_geometry(
     connect_edges: bool = False,
     interpolation_config: InterpolationConfig | None = None,
 ) -> tuple[list, list, list, PreprocessedData]:
-    data = _preprocess_data(DataManager().get_chart_data().parsed_data, data_type)
+    chart_data = DataManager().get_chart_data()
+    data = _preprocess_data(chart_data.parsed_data, data_type)
+    data.axis_labels = chart_data.labels
     verts = []
     edges = []
     faces = []