Make BaseScreen.start() and stop() idempotent.

7 files changed, 40 insertions, 42 deletions

diff --git a/urwid/curses_display.py b/urwid/curses_display.py
index 0a6155a..441042e 100755
--- a/urwid/curses_display.py
+++ b/urwid/curses_display.py
@@ -102,12 +102,10 @@ class Screen(BaseScreen, RealTerminal):
 
         self._mouse_tracking_enabled = enable
 
-    def start(self):
+    def _start(self):
         """
         Initialize the screen and input mode.
         """
-        assert self._started == False
-
         self.s = curses.initscr()
         self.has_color = curses.has_colors()
         if self.has_color:
@@ -130,15 +128,12 @@ class Screen(BaseScreen, RealTerminal):
         if not self._signal_keys_set:
             self._old_signal_keys = self.tty_signal_keys()
 
-        return super(Screen, self).start()
-
+        super(Screen, self)._start()
 
-    def stop(self):
+    def _stop(self):
         """
         Restore the screen.
         """
-        if self._started == False:
-            return
         curses.echo()
         self._curs_set(1)
         try:
@@ -149,7 +144,7 @@ class Screen(BaseScreen, RealTerminal):
         if self._old_signal_keys:
             self.tty_signal_keys(*self._old_signal_keys)
 
-        super(Screen, self).stop()
+        super(Screen, self)._stop()
 
 
     def _setup_colour_pairs(self):
diff --git a/urwid/display_common.py b/urwid/display_common.py
index 3f0e975..bf292b3 100755
--- a/urwid/display_common.py
+++ b/urwid/display_common.py
@@ -718,21 +718,35 @@ class BaseScreen(object):
 
     started = property(lambda self: self._started)
 
-    def start(self):
-        """Set up the screen.
+    def start(self, *args, **kwargs):
+        """Set up the screen.  If the screen has already been started, does
+        nothing.
 
-        May be used as a context manager, in which case `stop` will
+        May be used as a context manager, in which case :meth:`stop` will
         automatically be called at the end of the block:
 
             with screen.start():
                 ...
+
+        You shouldn't override this method in a subclass; instead, override
+        :meth:`_start`.
         """
+        if not self._started:
+            self._start(*args, **kwargs)
         self._started = True
         return StoppingContext(self)
 
+    def _start(self):
+        pass
+
     def stop(self):
+        if self._started:
+            self._stop()
         self._started = False
 
+    def _stop(self):
+        pass
+
     def run_wrapper(self, fn, *args, **kwargs):
         """Start the screen, call a function, then stop the screen.  Extra
         arguments are passed to `start`.
diff --git a/urwid/html_fragment.py b/urwid/html_fragment.py
index 86c7781..380d1d3 100755
--- a/urwid/html_fragment.py
+++ b/urwid/html_fragment.py
@@ -70,12 +70,6 @@ class HtmlGenerator(BaseScreen):
         """Not yet implemented"""
         pass
 
-    def start(self):
-        pass
-
-    def stop(self):
-        pass
-
     def set_input_timeouts(self, *args):
         pass
 
diff --git a/urwid/lcd_display.py b/urwid/lcd_display.py
index 13190bd..4f62173 100644
--- a/urwid/lcd_display.py
+++ b/urwid/lcd_display.py
@@ -33,12 +33,6 @@ class LCDScreen(BaseScreen):
     def set_mouse_tracking(self, enable=True):
         pass
 
-    def start(self):
-        pass
-
-    def stop(self):
-        pass
-
     def set_input_timeouts(self, *args):
         pass
 
diff --git a/urwid/main_loop.py b/urwid/main_loop.py
index f5f03dc..967596d 100755
--- a/urwid/main_loop.py
+++ b/urwid/main_loop.py
@@ -270,18 +270,13 @@ class MainLoop(object):
         Start the main loop handling input events and updating the screen. The
         loop will continue until an :exc:`ExitMainLoop` exception is raised.
 
-        This method will use :attr:`screen`'s run_wrapper() method if
-        :attr:`screen`'s start() method has not already been called.
-
         If you would prefer to manage the event loop yourself, don't use this
         method.  Instead, call :meth:`start` before starting the event loop,
         and :meth:`stop` once it's finished.
         """
         try:
-            if self.screen.started:
+            with self.screen.start():
                 self._run()
-            else:
-                self.screen.run_wrapper(self._run)
         except ExitMainLoop:
             pass
 
@@ -317,15 +312,20 @@ class MainLoop(object):
     def start(self):
         """
         Sets up the main loop, hooking into the event loop where necessary.
+        Starts the :attr:`screen` if it hasn't already been started.
 
         If you want to control starting and stopping the event loop yourself,
         you should call this method before starting, and call `stop` once the
-        loop has finished.
+        loop has finished.  You may also use this method as a context manager,
+        which will stop the loop automatically at the end of the block:
+
+            with main_loop.start():
+                ...
 
         Note that some event loop implementations don't handle exceptions
         specially if you manage the event loop yourself.  In particular, the
         Twisted and asyncio loops won't stop automatically when
-        :exc:`ExitMainLoop` is raised.
+        :exc:`ExitMainLoop` (or anything else) is raised.
         """
         if self.handle_mouse:
             self.screen.set_mouse_tracking()
diff --git a/urwid/raw_display.py b/urwid/raw_display.py
index e4cab86..6dbeb3f 100644
--- a/urwid/raw_display.py
+++ b/urwid/raw_display.py
@@ -189,13 +189,12 @@ class Screen(BaseScreen, RealTerminal):
         os.waitpid(self.gpm_mev.pid, 0)
         self.gpm_mev = None
 
-    def start(self, alternate_buffer=True):
+    def _start(self, alternate_buffer=True):
         """
         Initialize the screen and input mode.
 
         alternate_buffer -- use alternate screen buffer
         """
-        assert not self._started
         if alternate_buffer:
             self._term_output_file.write(escape.SWITCH_TO_ALTERNATE_BUFFER)
             self._rows_used = None
@@ -214,19 +213,17 @@ class Screen(BaseScreen, RealTerminal):
         if not self._signal_keys_set:
             self._old_signal_keys = self.tty_signal_keys(fileno=fd)
 
-        super(Screen, self).start()
-
         signals.emit_signal(self, INPUT_DESCRIPTORS_CHANGED)
         # restore mouse tracking to previous state
         self._mouse_tracking(self._mouse_tracking_enabled)
 
-    def stop(self):
+        return super(Screen, self)._start()
+
+    def _stop(self):
         """
         Restore the screen.
         """
         self.clear()
-        if not self._started:
-            return
 
         signals.emit_signal(self, INPUT_DESCRIPTORS_CHANGED)
 
@@ -254,7 +251,8 @@ class Screen(BaseScreen, RealTerminal):
         if self._old_signal_keys:
             self.tty_signal_keys(*(self._old_signal_keys + (fd,)))
 
-        super(Screen, self).stop()
+        super(Screen, self)._stop()
+
 
 
     def get_input(self, raw_keys=False):
diff --git a/urwid/web_display.py b/urwid/web_display.py
index 2258718..44a505c 100755
--- a/urwid/web_display.py
+++ b/urwid/web_display.py
@@ -632,7 +632,8 @@ class Screen:
         """
         global _prefs
 
-        assert not self._started
+        if self._started:
+            return util.StoppingContext(self)
 
         client_init = sys.stdin.read(50)
         assert client_init.startswith("window resize "),client_init
@@ -685,7 +686,9 @@ class Screen:
         """
         Restore settings and clean up.
         """
-        assert self._started
+        if not self._started:
+            return
+
         # XXX which exceptions does this actually raise? EnvironmentError?
         try:
             self._close_connection()