Move ready_to_plot

dstansby · dstansby · commit f2d899548c92 · 2023-06-13T13:43:31.000+01:00
diff --git a/src/napari_matplotlib/features.py b/src/napari_matplotlib/features.py
@@ -70,6 +70,22 @@ def _get_valid_axis_keys(self) -> List[str]:
         else:
             return self.layers[0].features.keys()
 
+    def _ready_to_plot(self) -> bool:
+        """
+        Return True if selected layer has a feature table we can plot with,
+        and the columns to plot have been selected.
+        """
+        if not hasattr(self.layers[0], "features"):
+            return False
+
+        feature_table = self.layers[0].features
+        valid_keys = self._get_valid_axis_keys()
+        return (
+            feature_table is not None
+            and len(feature_table) > 0
+            and all([self.get_key(dim) in valid_keys for dim in self.dims])
+        )
+
     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
@@ -114,28 +114,11 @@ def __init__(
         FeaturesMixin.__init__(self, ndim=2)
         self._update_layers(None)
 
-    def _ready_to_scatter(self) -> bool:
-        """
-        Return True if selected layer has a feature table we can scatter with,
-        and the two columns to be scatterd have been selected.
-        """
-        if not hasattr(self.layers[0], "features"):
-            return False
-
-        feature_table = self.layers[0].features
-        valid_keys = self._get_valid_axis_keys()
-        return (
-            feature_table is not None
-            and len(feature_table) > 0
-            and self.get_key("x") in valid_keys
-            and self.get_key("y") in valid_keys
-        )
-
     def draw(self) -> None:
         """
         Scatter two features from the currently selected layer.
         """
-        if self._ready_to_scatter():
+        if self._ready_to_plot():
             super().draw()
 
     def _get_data(self) -> Tuple[npt.NDArray[Any], npt.NDArray[Any], str, str]:
