Move data getting

dstansby · dstansby · commit ec8eabfa079d · 2023-06-13T13:56:49.000+01:00
diff --git a/src/napari_matplotlib/features.py b/src/napari_matplotlib/features.py
@@ -1,7 +1,10 @@
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 import napari
 import napari.layers
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
 from qtpy.QtWidgets import QComboBox, QLabel, QVBoxLayout
 
 from napari_matplotlib.base import NapariMPLWidget
@@ -86,6 +89,27 @@ def _ready_to_plot(self) -> bool:
             and all([self.get_key(dim) in valid_keys for dim in self.dims])
         )
 
+    def _get_data_names(
+        self,
+    ) -> Tuple[List[npt.NDArray[Any]], List[str]]:
+        """
+        Get the plot data from the ``features`` attribute of the first
+        selected layer.
+
+        Returns
+        -------
+        data : List[np.ndarray]
+            List contains X and Y columns from the FeatureTable. Returns
+            an empty array if nothing to plot.
+        names : List[str]
+            Names for each axis.
+        """
+        feature_table: pd.DataFrame = self.layers[0].features
+
+        names = [str(self.get_key(dim)) for dim in self.dims]
+        data = [np.array(feature_table[key]) for key in names]
+        return data, names
+
     def on_update_layers(self) -> None:
         """
         Called when the layer selection changes by ``self.update_layers()``.
diff --git a/src/napari_matplotlib/scatter.py b/src/napari_matplotlib/scatter.py
@@ -122,28 +122,5 @@ def draw(self) -> None:
             super().draw()
 
     def _get_data(self) -> Tuple[npt.NDArray[Any], npt.NDArray[Any], str, str]:
-        """
-        Get the plot data from the ``features`` attribute of the first
-        selected layer.
-
-        Returns
-        -------
-        data : List[np.ndarray]
-            List contains X and Y columns from the FeatureTable. Returns
-            an empty array if nothing to plot.
-        x_axis_name : str
-            The title to display on the x axis. Returns
-            an empty string if nothing to plot.
-        y_axis_name: str
-            The title to display on the y axis. Returns
-            an empty string if nothing to plot.
-        """
-        feature_table = self.layers[0].features
-
-        x = feature_table[self.get_key("x")]
-        y = feature_table[self.get_key("y")]
-
-        x_axis_name = str(self.get_key("x"))
-        y_axis_name = str(self.get_key("y"))
-
-        return x, y, x_axis_name, y_axis_name
+        data, names = self._get_data_names()
+        return data[0], data[1], names[0], names[1]
