Put run_wrapper in the base class; make BaseScreen.start() a contextmanager.

8 files changed, 39 insertions, 37 deletions

diff --git a/urwid/curses_display.py b/urwid/curses_display.py
index 758d621..0a6155a 100755
--- a/urwid/curses_display.py
+++ b/urwid/curses_display.py
@@ -130,7 +130,7 @@ class Screen(BaseScreen, RealTerminal):
         if not self._signal_keys_set:
             self._old_signal_keys = self.tty_signal_keys()
 
-        super(Screen, self).start()
+        return super(Screen, self).start()
 
 
     def stop(self):
@@ -152,19 +152,6 @@ class Screen(BaseScreen, RealTerminal):
         super(Screen, self).stop()
 
 
-    def run_wrapper(self,fn):
-        """Call fn in fullscreen mode.  Return to normal on exit.
-
-        This function should be called to wrap your main program loop.
-        Exception tracebacks will be displayed in normal mode.
-        """
-
-        try:
-            self.start()
-            return fn()
-        finally:
-            self.stop()
-
     def _setup_colour_pairs(self):
         """
         Initialize all 63 color pairs based on the term:
diff --git a/urwid/display_common.py b/urwid/display_common.py
index b539f53..3f0e975 100755
--- a/urwid/display_common.py
+++ b/urwid/display_common.py
@@ -26,7 +26,7 @@ try:
 except ImportError:
     pass # windows
 
-from urwid.util import int_scale
+from urwid.util import StoppingContext, int_scale
 from urwid import signals
 from urwid.compat import B, bytes3
 
@@ -719,11 +719,29 @@ class BaseScreen(object):
     started = property(lambda self: self._started)
 
     def start(self):
+        """Set up the screen.
+
+        May be used as a context manager, in which case `stop` will
+        automatically be called at the end of the block:
+
+            with screen.start():
+                ...
+        """
         self._started = True
+        return StoppingContext(self)
 
     def stop(self):
         self._started = False
 
+    def run_wrapper(self, fn, *args, **kwargs):
+        """Start the screen, call a function, then stop the screen.  Extra
+        arguments are passed to `start`.
+
+        Deprecated in favor of calling `start` as a context manager.
+        """
+        with self.start(*args, **kwargs):
+            return fn()
+
 
     def register_palette(self, palette):
         """Register a set of palette entries.
diff --git a/urwid/html_fragment.py b/urwid/html_fragment.py
index 159bc81..86c7781 100755
--- a/urwid/html_fragment.py
+++ b/urwid/html_fragment.py
@@ -82,10 +82,6 @@ class HtmlGenerator(BaseScreen):
     def reset_default_terminal_palette(self, *args):
         pass
 
-    def run_wrapper(self,fn):
-        """Call fn."""
-        return fn()
-
     def draw_screen(self, (cols, rows), r ):
         """Create an html fragment from the render object.
         Append it to HtmlGenerator.fragments list.
diff --git a/urwid/lcd_display.py b/urwid/lcd_display.py
index ca6336f..13190bd 100644
--- a/urwid/lcd_display.py
+++ b/urwid/lcd_display.py
@@ -45,9 +45,6 @@ class LCDScreen(BaseScreen):
     def reset_default_terminal_palette(self, *args):
         pass
 
-    def run_wrapper(self,fn):
-        return fn()
-
     def draw_screen(self, (cols, rows), r ):
         pass
 
diff --git a/urwid/main_loop.py b/urwid/main_loop.py
index da57f0d..f5f03dc 100755
--- a/urwid/main_loop.py
+++ b/urwid/main_loop.py
@@ -34,7 +34,7 @@ try:
 except ImportError:
     pass # windows
 
-from urwid.util import is_mouse_event
+from urwid.util import StoppingContext, is_mouse_event
 from urwid.compat import PYTHON3
 from urwid.command_map import command_map, REDRAW_SCREEN
 from urwid.wimp import PopUpTarget
@@ -343,6 +343,8 @@ class MainLoop(object):
         self._reset_input_descriptors()
         self.idle_handle = self.event_loop.enter_idle(self.entering_idle)
 
+        return StoppingContext(self)
+
     def stop(self):
         """
         Cleans up any hooks added to the event loop.  Only call this if you're
diff --git a/urwid/raw_display.py b/urwid/raw_display.py
index 000e4af..e4cab86 100644
--- a/urwid/raw_display.py
+++ b/urwid/raw_display.py
@@ -257,20 +257,6 @@ class Screen(BaseScreen, RealTerminal):
         super(Screen, self).stop()
 
 
-    def run_wrapper(self, fn, alternate_buffer=True):
-        """
-        Call start to initialize screen, then call fn.
-        When fn exits call stop to restore the screen to normal.
-
-        alternate_buffer -- use alternate screen buffer and restore
-            normal screen buffer on exit
-        """
-        try:
-            self.start(alternate_buffer)
-            return fn()
-        finally:
-            self.stop()
-
     def get_input(self, raw_keys=False):
         """Return pending input as a list.
 
diff --git a/urwid/util.py b/urwid/util.py
index f04e152..61e5eb6 100644
--- a/urwid/util.py
+++ b/urwid/util.py
@@ -458,3 +458,17 @@ def int_scale(val, val_range, out_range):
     # if num % dem == 0 then we are exactly half-way and have rounded up.
     return num // dem
 
+
+class StoppingContext(object):
+    """Context manager that calls ``stop`` on a given object on exit.  Used to
+    make the ``start`` method on `MainLoop` and `BaseScreen` optionally act as
+    context managers.
+    """
+    def __init__(self, wrapped):
+        self._wrapped = wrapped
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc_info):
+        self._wrapped.stop()
diff --git a/urwid/web_display.py b/urwid/web_display.py
index eccbad8..2258718 100755
--- a/urwid/web_display.py
+++ b/urwid/web_display.py
@@ -679,6 +679,8 @@ class Screen:
         signal.alarm( ALARM_DELAY )
         self._started = True
 
+        return util.StoppingContext(self)
+
     def stop(self):
         """
         Restore settings and clean up.