From 5c4b14468c9679f315f14432c341d540e65c77fe Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Tue, 24 Apr 2018 10:33:18 +0530
Subject: [PATCH] Apply default colors in diff kitten

---
 kittens/diff/main.py      | 14 +++++++++-----
 kittens/tui/handler.py    |  1 +
 kittens/tui/loop.py       |  6 +++++-
 kittens/tui/operations.py |  6 +++---
 4 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/kittens/diff/main.py b/kittens/diff/main.py
index 100f41a2b..9a5b02d24 100644
--- a/kittens/diff/main.py
+++ b/kittens/diff/main.py
@@ -12,7 +12,9 @@ from kitty.key_encoding import ESCAPE
 
 from ..tui.handler import Handler
 from ..tui.loop import Loop
-from ..tui.operations import clear_screen, set_line_wrapping, set_window_title
+from ..tui.operations import (
+    clear_screen, set_default_colors, set_line_wrapping, set_window_title
+)
 from .collect import create_collection, data_for_path
 from .config import init_config
 from .git import Differ
@@ -35,8 +37,9 @@ def generate_diff(collection, context):
 
 class DiffHandler(Handler):
 
-    def __init__(self, args, left, right):
+    def __init__(self, args, opts, left, right):
         self.state = INITIALIZING
+        self.opts = opts
         self.left, self.right = left, right
         self.report_traceback_on_exit = None
         self.args = args
@@ -54,6 +57,7 @@ class DiffHandler(Handler):
     def init_terminal_state(self):
         self.write(set_line_wrapping(False))
         self.write(set_window_title('kitty +diff'))
+        self.write(set_default_colors(self.opts.foreground, self.opts.background))
 
     def initialize(self):
         self.init_terminal_state()
@@ -61,7 +65,7 @@ class DiffHandler(Handler):
         self.create_collection()
 
     def finalize(self):
-        pass
+        self.write(set_default_colors())
 
     def draw_screen(self):
         if self.state < DIFFED:
@@ -144,10 +148,10 @@ def main(args):
     left, right = items
     if os.path.isdir(left) != os.path.isdir(right):
         raise SystemExit('The items to be diffed should both be either directories or files. Comparing a directory to a file is not valid.')
-    init_config(args)
+    opts = init_config(args)
 
     loop = Loop()
-    handler = DiffHandler(args, left, right)
+    handler = DiffHandler(args, opts, left, right)
     loop.loop(handler)
     if loop.return_code != 0:
         if handler.report_traceback_on_exit:
diff --git a/kittens/tui/handler.py b/kittens/tui/handler.py
index cccecbf18..ff00e3324 100644
--- a/kittens/tui/handler.py
+++ b/kittens/tui/handler.py
@@ -14,6 +14,7 @@ class Handler:
         self.initialize()
 
     def __exit__(self, *a):
+        del self.write_buf[:]
         self.finalize()
 
     def initialize(self):
diff --git a/kittens/tui/loop.py b/kittens/tui/loop.py
index 5fdadc749..e07accd48 100644
--- a/kittens/tui/loop.py
+++ b/kittens/tui/loop.py
@@ -376,8 +376,12 @@ class Loop:
                 self.return_code = 1
                 keep_going = False
 
+            finalize_output = b''.join(handler.write_buf).decode('utf-8')
+
             if tb is not None:
-                self._report_error_loop(tb, term_manager)
+                self._report_error_loop(finalize_output + tb, term_manager)
+        if tb is None:
+            os.write(self.output_fd, finalize_output.encode('utf-8'))
 
     def _report_error_loop(self, tb, term_manager):
         select = self.sel.select
diff --git a/kittens/tui/operations.py b/kittens/tui/operations.py
index c6cd0a6f1..e2bd446af 100644
--- a/kittens/tui/operations.py
+++ b/kittens/tui/operations.py
@@ -5,7 +5,7 @@
 import sys
 from contextlib import contextmanager
 
-from kitty.rgb import color_as_sharp, to_color
+from kitty.rgb import Color, color_as_sharp, to_color
 from kitty.terminfo import string_capabilities
 
 S7C1T = '\033 F'
@@ -168,9 +168,9 @@ def set_default_colors(fg=None, bg=None):
     if fg is None:
         ans += '\x1b]110\x1b\\'
     else:
-        ans += '\x1b]10;{}\x1b\\'.format(color_as_sharp(to_color(fg)))
+        ans += '\x1b]10;{}\x1b\\'.format(color_as_sharp(fg if isinstance(fg, Color) else to_color(fg)))
     if bg is None:
         ans += '\x1b]111\x1b\\'
     else:
-        ans += '\x1b]11;{}\x1b\\'.format(color_as_sharp(to_color(bg)))
+        ans += '\x1b]11;{}\x1b\\'.format(color_as_sharp(bg if isinstance(bg, Color) else to_color(bg)))
     return ans