ENH: In Chaco examples, replace gc.save_state()/gc.restore_state() wi…

…th 'with gc'. Remove unused imports.

examples/canvas/cliptest.py

@@ -3,8 +3,7 @@
 The main app for the PlotCanvas application
 """
 
-# Major library imports
-from numpy import arange, fabs, linspace, pi, sin
+from __future__ import with_statement
 
 # Enthought library imports
 from enthought.traits.api import Float
@@ -32,20 +31,19 @@ def __init__(self, *args, **kw):
         Component.__init__(self, *args, **kw)
 
     def _draw_mainlayer(self, gc, view_bounds=None, mode="default"):
-        gc.save_state()
-        gc.set_fill_color(self.fill_color)
-        dx, dy = self.bounds
-        x, y = self.position
-        gc.clip_to_rect(x, y, dx, dy)
-        gc.rect(x, y, dx, dy)
-        gc.fill_path()
-
-        ## draw line around outer box
-        #gc.set_stroke_color((0,0,0,1))
-        #gc.rect(self.outer_x, self.outer_y, self.outer_width, self.outer_height)
-        #gc.stroke_path()
+        with gc:
+            gc.set_fill_color(self.fill_color)
+            dx, dy = self.bounds
+            x, y = self.position
+            gc.clip_to_rect(x, y, dx, dy)
+            gc.rect(x, y, dx, dy)
+            gc.fill_path()
+            
+            ## draw line around outer box
+            #gc.set_stroke_color((0,0,0,1))
+            #gc.rect(self.outer_x, self.outer_y, self.outer_width, self.outer_height)
+            #gc.stroke_path()
 
-        gc.restore_state()
         return
 
     def normal_left_down(self, event):

examples/canvas/plot_clone_tool.py

@@ -1,10 +1,11 @@
 """ Makes a copy of the plot in the overlay and adds it to the canvas.
 """
 
+from __future__ import with_statement
 
 # Enthought library imports
 from enthought.traits.api import Bool, Callable, Enum, Float, Instance, Int, Trait, Tuple
-from enthought.enable.api import Canvas, Component, Container
+from enthought.enable.api import Container
 
 # Chaco imports
 from enthought.chaco.api import AbstractOverlay
@@ -50,14 +51,13 @@ def overlay(self, component, gc, view_bounds=None, mode="normal"):
         else:
             if self._offset is not None and (self._offset[0] > 10 or
                     self._offset[1] > 10):
-                gc.save_state()
-                gc.clear_clip_path()
-                gc.translate_ctm(*self._offset)
-                gc.set_alpha(self.alpha)
-                self._recursion_check = True
-                self.component._draw(gc, view_bounds, mode)
-                self._recursion_check = False
-                gc.restore_state()
+                with gc:
+                    gc.clear_clip_path()
+                    gc.translate_ctm(*self._offset)
+                    gc.set_alpha(self.alpha)
+                    self._recursion_check = True
+                    self.component._draw(gc, view_bounds, mode)
+                    self._recursion_check = False
 
     def drag_start(self, event):
         """ Called when the drag operation starts.  

examples/canvas/transient_plot_overlay.py

@@ -1,9 +1,11 @@
 
+from __future__ import with_statement
 
 from enthought.enable.api import Component
 from enthought.traits.api import Enum, Float, Instance, Trait, Tuple
 
-from enthought.chaco.api import AbstractOverlay, PlotComponent, BasePlotContainer
+from enthought.chaco.api import AbstractOverlay, BasePlotContainer
+
 
 class TransientPlotOverlay(BasePlotContainer, AbstractOverlay):
     """ Allows an arbitrary plot component to be overlaid on top of another one.
@@ -32,10 +34,9 @@ def _bounds_default(self):
 
     def overlay(self, component, gc, view_bounds=None, mode="normal"):
         self._do_layout()
-        gc.save_state()
-        gc.clear_clip_path()
-        self.overlay_component._draw(gc, view_bounds, mode)
-        gc.restore_state()
+        with gc:
+            gc.clear_clip_path()
+            self.overlay_component._draw(gc, view_bounds, mode)
 
     # TODO: Implement this more intelligently than the one in BasePlotContainer
     #def get_preferred_size(self):

examples/chaco_trait_editor.py

@@ -3,6 +3,8 @@
 interval.
 """
 
+from __future__ import with_statement
+
 from enthought.traits.ui.editor_factory import EditorFactory
 from enthought.traits.ui.wx.editor import Editor
 
@@ -17,6 +19,7 @@
 
 from math import pi
 
+
 class Interval(TraitType):
     """Trait that represents an interval.
 
@@ -39,6 +42,7 @@ def validate(self, object, name, value):
     def create_editor(self):
         return IntervalEditor()
 
+
 class IntervalEditorFactory(EditorFactory):
     width = Int(300)
     height = Int(40)
@@ -68,8 +72,7 @@ def overlay(self, component, gc, view_bounds=None, mode="normal"):
         coords = self._get_selection_screencoords()
         for coord in coords:
             start, end = coord
-            gc.save_state()
-            try:
+            with gc:
                 gc.set_alpha(self.alpha)
                 gc.set_stroke_color(self.border_color_)
                 gc.set_line_width(self.border_width)
@@ -87,14 +90,11 @@ def overlay(self, component, gc, view_bounds=None, mode="normal"):
                 gc.set_fill_color(self.high_color_)
                 self._circle(gc, end, mid_y, self.radius)
                 gc.draw_path()
-            finally:
-                gc.restore_state()
 
     def _circle(self, gc, x, y, radius):
-        gc.save_state()
-        gc.translate_ctm(x, y)
-        gc.arc(0, 0, 2*radius, 0, 2*pi)
-        gc.restore_state()
+        with gc:
+            gc.translate_ctm(x, y)
+            gc.arc(0, 0, 2*radius, 0, 2*pi)
 
 
 class IntervalEditorImpl(Editor):
@@ -159,7 +159,7 @@ def update_editor(self):
 # --- Demonstration ---
 
 if __name__ == "__main__":
-    from enthought.traits.api import HasTraits, Range
+    from enthought.traits.api import HasTraits
     from enthought.traits.ui.api import View, Item
     class IntervalTest(HasTraits):
         interval = Interval(low=0, high=1)

examples/logo.py

@@ -1,11 +1,14 @@
 """ LOGO overlay """
 
+from __future__ import with_statement
+
 from numpy import array, cos, invert, isnan, nan, pi, sin, vstack
-from enthought.traits.api import Array, Enum, Float, HasTraits, Instance, Range
+from enthought.traits.api import Array, Enum, Float, Range
 from enthought.traits.ui.api import Group, Item, View
 from enthought.enable.api import ColorTrait
 from enthought.chaco.api import arg_find_runs, AbstractOverlay
 
+
 class Turtle(AbstractOverlay):
     x = Float
     y = Float
@@ -30,8 +33,7 @@ def overlay(self, other_component, gc, view_bounds=None, mode="normal"):
         self.render(gc, other_component)
 
     def render_turtle(self, gc, component):
-        gc.save_state()
-        try:
+        with gc:
             x, y = component.map_screen(array([self.x, self.y], ndmin=2))[0]
             gc.translate_ctm(x, y)
             angle = self.angle * pi / 180.0
@@ -43,25 +45,20 @@ def render_turtle(self, gc, component):
                       [-0.707*self.size, -0.707*self.size],
                       [self.size, 0.0]])
             gc.fill_path()
-        finally:
-            gc.restore_state()
 
     def render(self, gc, component):
         # Uses the component to map our path into screen space
         nan_mask = invert(isnan(self.path[:,0])).astype(int)
         blocks = [b for b in arg_find_runs(nan_mask, "flat") if nan_mask[b[0]] != 0]
         screen_pts = component.map_screen(self.path)
-        gc.save_state()
-        try:
+        with gc:
             gc.clip_to_rect(component.x, component.y, component.width, component.height)
             gc.set_stroke_color(self.line_color_)
             for start, end in blocks:
                 gc.begin_path()
                 gc.lines(screen_pts[start:end])
                 gc.stroke_path()
             self.render_turtle(gc, component)
-        finally:
-            gc.restore_state()
 
     def pendown(self):
         self._pen = "down"

examples/xray_plot.py

@@ -2,6 +2,8 @@
 Implementation of a plut using a custom overlay and tool
 """
 
+from __future__ import with_statement
+
 import numpy
 
 from enthought.traits.api import HasTraits, Instance, Enum
@@ -77,17 +79,11 @@ def overlay(self, component, gc, view_bounds=None, mode='normal'):
         x1, x2 = x_range
         y1, y2 = y_range
 
-        # whenever save_state is called, always make sure restore_state is 
-        # called
-
-        gc.save_state()
-        try:
+        with gc:
             gc.set_alpha(0.8)
             gc.set_fill_color((1.0,1.0,1.0))
             gc.rect(x1, y1, x2-x1, y2-y1)
             gc.draw_path()
-        finally:
-            gc.restore_state()      
 
         pts = self._get_selected_points()
         if len(pts) == 0:

examples/zoomed_plot/zoom_overlay.py

@@ -1,5 +1,7 @@
 
-from numpy import array, amax, amin, clip
+from __future__ import with_statement
+
+from numpy import array, amax, amin
 
 from enthought.enable.api import ColorTrait, Component
 from enthought.traits.api import Float, Instance, Int
@@ -64,8 +66,7 @@ def overlay(self, component, gc, view_bounds=None, mode="normal"):
 
         left_line, right_line, polygon = self.calculate_points(component)
 
-        gc.save_state()
-        try:
+        with gc:
             gc.translate_ctm(*component.position)
             gc.set_alpha(self.alpha)
             gc.set_fill_color(self.fill_color_)
@@ -79,8 +80,7 @@ def overlay(self, component, gc, view_bounds=None, mode="normal"):
             gc.lines(left_line)
             gc.lines(right_line)
             gc.stroke_path()
-        finally:
-            gc.restore_state()
+
         return
 
     def _get_selection_screencoords(self):