Add MappedUnion and Optional traits and use Optional where appropriate

chaco/base_candle_plot.py

@@ -17,6 +17,7 @@
 
 # Chaco imports
 from .base_xy_plot import BaseXYPlot
+from .chaco_traits import Optional
 
 # TODO: allow to set the width of the bar
 
@@ -54,11 +55,11 @@ class BaseCandlePlot(BaseXYPlot):
     #: The color of the stems reaching from the bar ends to the min and max
     #: values.  Also the color of the endcap line segments at min and max.  If
     #: None, this defaults to **bar_line_color**.
-    stem_color = Union(None, ColorTrait("black"))
+    stem_color = Optional(ColorTrait("black"))
 
     #: The color of the line drawn across the bar at the center values.
     #: If None, this defaults to **bar_line_color**.
-    center_color = Union(None, ColorTrait("black"))
+    center_color = Optional(ColorTrait("black"))
 
     #: The color of the outline to draw around the bar.
     outline_color = ColorTrait("black")
@@ -69,11 +70,11 @@ class BaseCandlePlot(BaseXYPlot):
 
     #: The thickness, in pixels, of the stem lines.  If None, this defaults
     #: to **line_width**.
-    stem_width = Union(None, Int(1))
+    stem_width = Optional(Int(1))
 
     #: The thickeness, in pixels, of the line drawn across the bar at the
     #: center values.  If None, this defaults to **line_width**.
-    center_width = Union(None, Int(1))
+    center_width = Optional(Int(1))
 
     #: Whether or not to draw bars at the min and max extents of the error bar
     end_cap = Bool(True)

chaco/chaco_traits.py

@@ -12,7 +12,7 @@
 """
 
 # Enthought library imports
-from traits.api import Enum
+from traits.api import Enum, Union, TraitError
 
 # ----------------------------------------------------------------------------
 # Box positioning traits: used to specify positions of boxes relative to
@@ -24,3 +24,59 @@
 #: Values correspond to: top, bottom, left, right, top left, top right, bottom
 #: left, bottom right
 box_position_enum = Enum("T", "B", "L", "R", "TL", "TR", "BL", "BR")
+
+
+class MappedUnion(Union):
+    """Version of the Union trait that handles mapped traits correctly."""
+
+    #: This is not mapped by default.
+    is_mapped = False
+
+    def __init__(self, *traits, **metadata):
+        super().__init__(*traits, **metadata)
+
+        # look for post_setattr and is_mapped on traits
+        post_setattrs = []
+        mapped_traits = []
+        for trait in traits:
+            if trait is None:
+                continue
+            post_setattr = getattr(trait, "post_setattr", None)
+            if post_setattr is not None:
+                post_setattrs.append(post_setattr)
+            if trait.is_mapped:
+                self.is_mapped = True
+                mapped_traits.append(trait)
+
+        if post_setattrs:
+            self.post_setattrs = post_setattrs
+            self.post_setattr = self._post_setattr
+        if self.is_mapped:
+            self.mapped_traits = mapped_traits
+
+    def mapped_value(self, value):
+        for trait in self.mapped_traits:
+            try:
+                return trait.mapped_value(value)
+            except Exception:
+                pass
+
+        return value
+
+    def _post_setattr(self, object, name, value):
+        for post_setattr in self.post_setattrs:
+            try:
+                post_setattr(object, name, value)
+                return
+            except Exception:
+                pass
+
+        if self.is_mapped:
+            setattr(object, name + "_", value)
+
+
+class Optional(MappedUnion):
+    """Convenience class"""
+
+    def __init__(self, trait, **metadata):
+        super().__init__(None, trait, **metadata)

chaco/data_view.py

@@ -20,6 +20,7 @@
 from .axis import PlotAxis
 from .base_1d_mapper import Base1DMapper
 from .base_2d_plot import Base2DPlot
+from .chaco_traits import Optional
 from .data_range_2d import DataRange2D
 from .grid import PlotGrid
 from .linear_mapper import LinearMapper
@@ -236,10 +237,10 @@ class DataView(OverlayPlotContainer):
         observe='y_axis.[title,orientation], x_axis.[title,orientation]'
     )
 
-    _padding_top = Union(None, Int())
-    _padding_bottom = Union(None, Int())
-    _padding_left = Union(None, Int())
-    _padding_right = Union(None, Int())
+    _padding_top = Optional(Int())
+    _padding_bottom = Optional(Int())
+    _padding_left = Optional(Int())
+    _padding_right = Optional(Int())
 
     def _find_padding(self, side):
         SIDE_TO_TRAIT_MAP = {

chaco/grid.py

@@ -46,6 +46,7 @@
 # Local, relative imports
 from .abstract_overlay import AbstractOverlay
 from .abstract_mapper import AbstractMapper
+from .chaco_traits import Optional
 from .log_mapper import LogMapper
 from .ticks import AbstractTickGenerator, DefaultTickGenerator
 
@@ -110,11 +111,11 @@ class PlotGrid(AbstractOverlay):
 
     #: The dataspace value at which to start this grid.  If None, then
     #: uses the mapper.range.low.
-    data_min = Union(None, Float)
+    data_min = Optional(Float)
 
     #: The dataspace value at which to end this grid.  If None, then uses
     #: the mapper.range.high.
-    data_max = Union(None, Float)
+    data_max = Optional(Float)
 
     #: A callable that implements the AbstractTickGenerator Interface.
     tick_generator = Instance(AbstractTickGenerator)

chaco/overlays/scatter_inspector_overlay.py

@@ -13,11 +13,12 @@
 
 # Enthought library imports
 from enable.api import ColorTrait, MarkerTrait
-from traits.api import Float, Int, Str, Union
+from traits.api import Float, Int, Str
 from traits.observation.events import TraitChangeEvent
 
 # Local, relative imports
 from chaco.abstract_overlay import AbstractOverlay
+from chaco.chaco_traits import Optional
 from chaco.plots.scatterplot import render_markers
 
 
@@ -32,19 +33,19 @@ class ScatterInspectorOverlay(AbstractOverlay):
 
     #: The style to use when a point is hovered over
     hover_metadata_name = Str("hover")
-    hover_marker = Union(None, MarkerTrait)
-    hover_marker_size = Union(None, Int)
-    hover_line_width = Union(None, Float)
-    hover_color = Union(None, ColorTrait)
-    hover_outline_color = Union(None, ColorTrait)
+    hover_marker = Optional(MarkerTrait)
+    hover_marker_size = Optional(Int)
+    hover_line_width = Optional(Float)
+    hover_color = Optional(ColorTrait)
+    hover_outline_color = Optional(ColorTrait)
 
     #: The style to use when a point has been selected by a click
     selection_metadata_name = Str("selections")
-    selection_marker = Union(None, MarkerTrait)
-    selection_marker_size = Union(None, Int)
-    selection_line_width = Union(None, Float)
-    selection_color = Union(None, ColorTrait)
-    selection_outline_color = Union(None, ColorTrait)
+    selection_marker = Optional(MarkerTrait)
+    selection_marker_size = Optional(Int)
+    selection_line_width = Optional(Float)
+    selection_color = Optional(ColorTrait)
+    selection_outline_color = Optional(ColorTrait)
 
     # For now, implement the equivalent of this Traits 3 feature manually
     # using a series of trait change handlers (defined at the end of the

chaco/plots/multi_line_plot.py

@@ -39,6 +39,7 @@
 from chaco.array_data_source import ArrayDataSource
 from chaco.base import arg_find_runs, bin_search
 from chaco.base_xy_plot import BaseXYPlot
+from chaco.chaco_traits import Optional
 
 
 class MultiLinePlot(BaseXYPlot):
@@ -118,7 +119,7 @@ class MultiLinePlot(BaseXYPlot):
     color = black_color_trait(requires_redraw=True)
 
     #: A function that returns the color of lines.  Overrides `color` if not None.
-    color_func = Union(None, Callable)
+    color_func = Optional(Callable)
 
     #: The color to use to highlight the line when selected.
     selected_color = ColorTrait("lightyellow")

chaco/plotscrollbar.py

@@ -12,6 +12,8 @@
 
 from enable.api import NativeScrollBar
 
+from .chaco_traits import Optional
+
 
 class PlotScrollBar(NativeScrollBar):
     """
@@ -43,7 +45,7 @@ class PlotScrollBar(NativeScrollBar):
     _mapper = Any()
 
     # Stores the index (0 or 1) corresponding to self.axis
-    _axis_index = Union(None, Int)
+    _axis_index = Optional(Int)
 
     # ----------------------------------------------------------------------
     # Public methods

chaco/tests/test_chaco_traits.py

@@ -0,0 +1,117 @@
+# (C) Copyright 2005-2021 Enthought, Inc., Austin, TX
+# All rights reserved.
+#
+# This software is provided without warranty under the terms of the BSD
+# license included in LICENSE.txt and may be redistributed only under
+# the conditions described in the aforementioned license. The license
+# is also available online at http://www.enthought.com/licenses/BSD.txt
+#
+# Thanks for using Enthought open source!
+
+import unittest
+
+from traits.api import Float, HasTraits, Int, Map, Str
+from enable.api import ColorTrait
+
+from ..chaco_traits import MappedUnion, Optional
+
+
+class UsesMappedUnion(HasTraits):
+
+    no_mapped = MappedUnion(Int, Float)
+
+    mapped = MappedUnion(
+        None, Str, Map({'yes': True, 'no': False}), ColorTrait
+    )
+
+    optional = Optional(ColorTrait)
+
+    optional_default = Optional(ColorTrait, default_value='red')
+
+
+class TestMappedUnion(unittest.TestCase):
+
+    def test_no_mapped(self):
+        no_mapped = MappedUnion(Int, Float)
+
+        self.assertFalse(no_mapped.is_mapped)
+        self.assertIsNone(no_mapped.post_setattr)
+
+    def test_mapped(self):
+        mapped = MappedUnion(None, Str, Map({'yes': True, 'no': False}))
+
+        self.assertTrue(mapped.is_mapped)
+        self.assertIsNotNone(mapped.post_setattr)
+
+    def test_optional(self):
+        mapped = Optional(ColorTrait)
+
+        self.assertTrue(mapped.is_mapped)
+        self.assertIsNotNone(mapped.post_setattr)
+
+    def test_no_mapped_class(self):
+        mapped_union = UsesMappedUnion()
+
+        no_mapped = mapped_union.trait('no_mapped')
+
+        self.assertFalse(no_mapped.is_mapped)
+        self.assertIsNone(no_mapped.post_setattr)
+
+        self.assertFalse(hasattr(mapped_union, 'no_mapped_'))
+
+        mapped_union.no_mapped = 1
+
+        self.assertFalse(hasattr(mapped_union, 'no_mapped_'))
+
+    def test_mapped_class(self):
+        mapped_union = UsesMappedUnion()
+
+        mapped = mapped_union.trait('mapped')
+
+        self.assertTrue(mapped.is_mapped)
+        self.assertIsNotNone(mapped.post_setattr)
+
+        # test default
+        self.assertIsNone(mapped_union.mapped_)
+
+        # test mapper works
+        mapped_union.mapped = 'yes'
+
+        self.assertTrue(mapped_union.mapped_)
+
+        # test second mapper works
+        mapped_union.mapped = 'red'
+
+        self.assertEqual(mapped_union.mapped_, (1.0, 0.0, 0.0, 1.0))
+
+        # test non-mapped value works
+        mapped_union.mapped = 'notacolor'
+
+        self.assertEqual(mapped_union.mapped_, 'notacolor')
+
+    def test_optional_class(self):
+        mapped_union = UsesMappedUnion()
+
+        optional = mapped_union.trait('optional')
+
+        self.assertTrue(optional.is_mapped)
+        self.assertIsNotNone(optional.post_setattr)
+
+        # test default
+        self.assertIsNone(mapped_union.optional_)
+
+        # test mapper works
+        mapped_union.optional = 'red'
+
+        self.assertEqual(mapped_union.optional_, (1.0, 0.0, 0.0, 1.0))
+
+        # test non-mapped value works
+        mapped_union.optional = None
+
+        self.assertIsNone(mapped_union.optional_)
+
+    def test_optional_default_class(self):
+        mapped_union = UsesMappedUnion()
+
+        # test default
+        self.assertEqual(mapped_union.optional_default_, (1.0, 0.0, 0.0, 1.0))

chaco/tools/better_selecting_zoom.py

@@ -10,7 +10,6 @@
 
 import numpy
 
-from chaco.abstract_overlay import AbstractOverlay
 from enable.api import ColorTrait, KeySpec
 from traits.api import (
     Bool,
@@ -22,6 +21,8 @@
     Union
 )
 
+from chaco.abstract_overlay import AbstractOverlay
+from chaco.chaco_traits import Optional
 from .better_zoom import BetterZoom
 from .tool_states import SelectedZoomState
 
@@ -91,10 +92,10 @@ class BetterSelectingZoom(AbstractOverlay, BetterZoom):
     event_state = Enum("normal", "selecting", "pre_selecting")
 
     # The (x,y) screen point where the mouse went down.
-    _screen_start = Union(None, Tuple)
+    _screen_start = Optional(Tuple)
 
     # The (x,,y) screen point of the last seen mouse move event.
-    _screen_end = Union(None, Tuple)
+    _screen_end = Optional(Tuple)
 
     # If **always_on** is False, this attribute indicates whether the tool
     # is currently enabled.

chaco/tools/image_inspector_tool.py

@@ -17,6 +17,7 @@
 
 # Chaco imports
 from chaco.abstract_overlay import AbstractOverlay
+from chaco.chaco_traits import Optional
 from chaco.overlays.text_box_overlay import TextBoxOverlay
 from chaco.plots.image_plot import ImagePlot
 
@@ -42,7 +43,7 @@ class ImageInspectorTool(BaseTool):
 
     # Stores the value of self.visible when the mouse leaves the tool,
     # so that it can be restored when the mouse enters again.
-    _old_visible = Union(None, Bool)
+    _old_visible = Optional(Bool)
 
     def normal_key_pressed(self, event):
         if self.inspector_key.match(event):