tui: Allow handlers to finalize() as well as initialize()

kittens/diff/main.py

@@ -54,8 +54,7 @@ class DiffHandler(Handler):
         self.write(set_line_wrapping(False))
         self.write(set_window_title('kitty +diff'))
 
-    def initialize(self, *args):
-        Handler.initialize(self, *args)
+    def initialize(self):
         self.init_terminal_state()
         self.draw_screen()
         self.create_collection()

kittens/hints/main.py

@@ -100,8 +100,7 @@ class Hints(Handler):
         self.write(set_cursor_visible(False))
         self.write(set_window_title(self.window_title))
 
-    def initialize(self, *args):
-        Handler.initialize(self, *args)
+    def initialize(self):
         self.init_terminal_state()
         self.draw_screen()
 
@@ -211,7 +210,7 @@ def run(args, text):
         if chars is None:
             import json
             chars = json.loads(os.environ['KITTY_COMMON_OPTS'])['select_by_word_characters']
-        pat = re.compile('(?u)[{}\w]{{{},}}'.format(escape(chars), args.minimum_match_length))
+        pat = re.compile(r'(?u)[{}\w]{{{},}}'.format(escape(chars), args.minimum_match_length))
         finditer = partial(regex_finditer, pat, args.minimum_match_length)
     else:
         finditer = partial(regex_finditer, re.compile(args.regex), args.minimum_match_length)

kittens/tui/handler.py

@@ -5,11 +5,23 @@
 
 class Handler:
 
-    def initialize(self, screen_size, quit_loop, wakeup, start_job):
+    def _initialize(self, screen_size, quit_loop, wakeup, start_job):
         self.screen_size, self.quit_loop = screen_size, quit_loop
         self.wakeup = wakeup
         self.start_job = start_job
 
+    def __enter__(self):
+        self.initialize()
+
+    def __exit__(self, *a):
+        self.finalize()
+
+    def initialize(self):
+        pass
+
+    def finalize(self):
+        pass
+
     def on_resize(self, screen_size):
         self.screen_size = screen_size
 

kittens/tui/loop.py

@@ -121,8 +121,7 @@ class UnhandledException(Handler):
     def __init__(self, tb):
         self.tb = tb
 
-    def initialize(self, *args):
-        Handler.initialize(self, *args)
+    def initialize(self):
         self.write(clear_screen())
         self.write(self.tb.replace('\n', '\r\n'))
         self.write('\r\n')
@@ -359,29 +358,24 @@ class Loop:
             handler._term_manager = term_manager
             keep_going = True
             try:
-                handler.initialize(screen_size(), self.quit, self.wakeup, self.start_job)
+                handler._initialize(screen_size(), self.quit, self.wakeup, self.start_job)
+                with handler:
+                    while keep_going:
+                        has_data_to_write = bool(handler.write_buf)
+                        if not has_data_to_write and not self.read_allowed:
+                            break
+                        if has_data_to_write != waiting_for_write:
+                            waiting_for_write = has_data_to_write
+                            self._modify_output_selector(waiting_for_write)
+                        events = select()
+                        for key, mask in events:
+                            key.data(handler)
             except Exception:
                 import traceback
                 tb = traceback.format_exc()
                 self.return_code = 1
                 keep_going = False
-            while keep_going:
-                has_data_to_write = bool(handler.write_buf)
-                if not has_data_to_write and not self.read_allowed:
-                    break
-                if has_data_to_write != waiting_for_write:
-                    waiting_for_write = has_data_to_write
-                    self._modify_output_selector(waiting_for_write)
-                events = select()
-                for key, mask in events:
-                    try:
-                        key.data(handler)
-                    except Exception:
-                        import traceback
-                        tb = traceback.format_exc()
-                        self.return_code = 1
-                        keep_going = False
-                        break
+
             if tb is not None:
                 self._report_error_loop(tb, term_manager)
 
@@ -391,14 +385,15 @@ class Loop:
         handler = UnhandledException(tb)
         handler.write_buf = []
         handler._term_manager = term_manager
-        handler.initialize(screen_size(), self.quit, self.wakeup, self.start_job)
-        while True:
-            has_data_to_write = bool(handler.write_buf)
-            if not has_data_to_write and not self.read_allowed:
-                break
-            if has_data_to_write != waiting_for_write:
-                waiting_for_write = has_data_to_write
-                self._modify_output_selector(waiting_for_write)
-            events = select()
-            for key, mask in events:
-                key.data(handler)
+        handler._initialize(screen_size(), self.quit, self.wakeup, self.start_job)
+        with handler:
+            while True:
+                has_data_to_write = bool(handler.write_buf)
+                if not has_data_to_write and not self.read_allowed:
+                    break
+                if has_data_to_write != waiting_for_write:
+                    waiting_for_write = has_data_to_write
+                    self._modify_output_selector(waiting_for_write)
+                events = select()
+                for key, mask in events:
+                    key.data(handler)

kittens/unicode_input/main.py

@@ -327,8 +327,7 @@ class UnicodeInput(Handler):
         self.write(set_line_wrapping(False))
         self.write(set_window_title(_('Unicode input')))
 
-    def initialize(self, *args):
-        Handler.initialize(self, *args)
+    def initialize(self):
         self.init_terminal_state()
         self.draw_screen()