Cleanup map_screen return types

chaco/data_view.py

@@ -242,7 +242,11 @@ def map_screen(self, data_array):
         """Maps an array of data points to screen space and returns an array
         of screen space points.
         """
-        # data_array is Nx2 array
+
+        # ensure data_array is an Nx2 ndarray
+        data_array = array(data_array)
+        data_array = data_array.reshape(-1,2)
+
         if len(data_array) == 0:
             return empty(shape=(0,2))
         x_ary, y_ary = transpose(data_array)

chaco/examples/demo/basic/hittest_tool.py

@@ -53,7 +53,7 @@ def normal_mouse_move(self, event):
     def overlay(self, plot, gc, view_bounds=None, mode="normal"):
         # If we have a point, draw it to the screen as a small square
         if self.pt is not None:
-            x, y = plot.map_screen(self.pt)
+            x, y = plot.map_screen(self.pt)[0]
             gc.draw_rect((int(x) - 2, int(y) - 2, 4, 4))
 
 

chaco/grid_mapper.py

@@ -17,7 +17,7 @@
 from contextlib import contextmanager
 
 # Major library imports
-from numpy import column_stack, transpose
+from numpy import array, column_stack, transpose
 
 # Enthought library imports
 from traits.api import Bool, DelegatesTo, Instance, Float, Property
@@ -126,6 +126,10 @@ def map_screen(self, data_pts):
 
         Maps values from data space into screen space.
         """
+        # ensure data_array is an Nx2 ndarray
+        data_pts = array(data_pts)
+        data_pts = data_pts.reshape(-1, 2)
+
         xs, ys = transpose(data_pts)
         screen_xs = self._xmapper.map_screen(xs)
         screen_ys = self._ymapper.map_screen(ys)

chaco/plots/horizon_plot.py

@@ -30,6 +30,8 @@ def map_screen(self, data_array):
                 return array([self.low_pos])
         else:
             # Scale the data by the number of bands
+            if not isinstance(data_array, ndarray):
+                data_array = array(data_array, ndmin=1)
             return (
                 data_array * self.bands - self.range.low
             ) * self._scale + self.low_pos

chaco/plots/tests/test_image_plot.py

@@ -186,6 +186,25 @@ def test_vertical_bottom_right(self):
             RGB, (IMAGE.T)[::-1, ::-1], origin="bottom right", orientation="v"
         )
 
+    def test_map_screen(self):
+        data_source = ImageData(data=RGB)
+        index, index_mapper = get_image_index_and_mapper(RGB)
+        renderer = ImagePlot(
+            value=data_source,
+            index=index,
+            index_mapper=index_mapper,
+        )
+
+        # ImagePlot map screen is used tto find the screen_bbox, not on data
+        # itself.
+        screen_pt = renderer.map_screen([(0, 0)])
+        self.assertEqual(type(screen_pt), np.ndarray)
+        self.assertEqual(screen_pt.shape, (1, 2))
+
+        screen_pt = renderer.map_screen([])
+        self.assertEqual(type(screen_pt), np.ndarray)
+        self.assertEqual(screen_pt.shape, (0, 2))
+
     # regression test for enthought/chaco#528
     @unittest.skipIf(is_null, "Skip on 'null' toolkit")
     def test_resize_to_zero(self):

chaco/polar_mapper.py

@@ -12,8 +12,6 @@
 Defines the PolarMapper class, which maps from a 1-D region in data space
 into a 1-D output space.
 """
-
-# Major library imports
 from numpy import array, float64, full_like, ndarray
 
 # Enthought library imports
@@ -62,6 +60,8 @@ def map_screen(self, data_array):
             else:
                 return array([self.low_pos])
         else:
+            if not isinstance(data_array, ndarray):
+                data_array = array(data_array, ndmin=1)
             return (data_array - self.range.low) * self._scale + self.low_pos
 
     def map_data(self, screen_val):

examples/tutorials/scipy2008/custom_overlay_dataspace.py

@@ -43,7 +43,7 @@ class CustomOverlay(AbstractOverlay):
 
     def overlay(self, component, gc, view_bounds=None, mode="normal"):
         if self.dataspace:
-            self.x, self.y = component.map_screen(self._anchor)
+            self.x, self.y = component.map_screen(self._anchor)[0]
         gc.set_fill_color(self.color_)
         x = self.x + component.x
         y = self.y + component.y

examples/tutorials/scipy2008/custom_overlay_movetool.py

@@ -53,7 +53,7 @@ class CustomOverlay(AbstractOverlay):
 
     def overlay(self, component, gc, view_bounds=None, mode="normal"):
         if self.dataspace:
-            self.x, self.y = component.map_screen(self._anchor)
+            self.x, self.y = component.map_screen(self._anchor)[0]
         gc.set_fill_color(self.color_)
         x = self.x + component.x
         y = self.y + component.y