Merge pull request #4 from mil-ad/basic-tmux-support

Adding tmux passthrough support

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kitcat"
-version = "1.1.1"
+version = "1.2.0"
 description = "Matplotlib backend for direct plotting in the terminal using Kitty graphics protocol."
 readme = "README.md"
 requires-python = ">=3.9"

src/kitcat/backend.py

@@ -2,26 +2,62 @@
 import os
 import sys
 from base64 import b64encode
+from functools import wraps
 from io import BytesIO
 
 from matplotlib.backend_bases import FigureManagerBase
 from matplotlib.backends.backend_agg import FigureCanvasAgg
 
+from .utils import num_required_lines
+
 __all__ = ["FigureCanvas", "FigureManager"]
 
 CHUNK_SIZE_KITTY = 4096
 CHUNK_SIZE_IT2 = 1_048_576
 
 
-def display_kitty(pixel_data):
+def make_tmux_compatible(func):
+    if "TMUX" not in os.environ:
+        return func
+
+    old_write_fn = sys.stdout.write
+
+    def new_write_fn(s):
+        s = s.replace("\033", "\033\033")
+        old_write_fn(s)
+
+    @wraps(func)
+    def wrapper(img_buf):
+        try:
+            height_lines = num_required_lines(img_buf)
+            sys.stdout.write("\n" * height_lines)
+            # sys.stdout.write("\033[?25l")
+            sys.stdout.write(f"\033[{height_lines}F")
+            sys.stdout.write("\033Ptmux;")
+
+            sys.stdout.write = new_write_fn
+            func(img_buf)
+            sys.stdout.write = old_write_fn
+
+            sys.stdout.write("\033\\")
+            sys.stdout.write(f"\033[{height_lines}E")
+        finally:
+            # Ensure stdout is always restored
+            sys.stdout.write = old_write_fn
+
+    return wrapper
+
+
+@make_tmux_compatible
+def display_kitty(img_buf):
     """
     Encodes pixel data to the terminal using Kitty graphics protocol. All escape codes
     are of the form: <ESC>_G<control data>;<payload><ESC>\
 
     For more information on the protocol see:
     https://sw.kovidgoyal.net/kitty/graphics-protocol/#control-data-reference
     """
-    data = b64encode(pixel_data).decode("ascii")
+    data = b64encode(img_buf.read()).decode("ascii")
 
     first_chunk, more_data = data[:CHUNK_SIZE_KITTY], data[CHUNK_SIZE_KITTY:]
 
@@ -36,9 +72,6 @@ def display_kitty(pixel_data):
         chunk, more_data = more_data[:CHUNK_SIZE_KITTY], more_data[CHUNK_SIZE_KITTY:]
         sys.stdout.write(f"\033_Gm={'1' if more_data else '0'};{chunk}\033\\")
 
-    sys.stdout.write("\n")
-    sys.stdout.flush()
-
 
 def display_iterm2_new(pixel_data):
     data = b64encode(pixel_data).decode("ascii")
@@ -51,13 +84,13 @@ def display_iterm2_new(pixel_data):
     sys.stdout.flush()
 
 
-def display_iterm2(pixel_data):
+@make_tmux_compatible
+def display_iterm2(img_buf):
+    pixel_data = img_buf.read()
     data = b64encode(pixel_data).decode("ascii")
 
     # size is optional in iTerm2 but is required in vscode terminal
     sys.stdout.write(f"\033]1337;File=inline=1;size={len(pixel_data)}:{data}\a")
-    sys.stdout.write("\n")
-    sys.stdout.flush()
 
 
 class KitcatFigureManager(FigureManagerBase):
@@ -67,9 +100,12 @@ def show(self):
             buf.seek(0)
 
             if os.environ.get("TERM_PROGRAM") in ["iTerm.app", "vscode"]:
-                display_iterm2(pixel_data=buf.read())
+                display_iterm2(img_buf=buf)
             else:
-                display_kitty(pixel_data=buf.read())
+                display_kitty(img_buf=buf)
+
+        sys.stdout.write("\n")
+        sys.stdout.flush()
 
 
 class KitcatFigureCanvas(FigureCanvasAgg):

src/kitcat/utils.py

@@ -0,0 +1,25 @@
+import array
+import fcntl
+import math
+import sys
+import termios
+
+from PIL import Image
+
+
+def get_char_cell_height() -> int:
+    """Source https://sw.kovidgoyal.net/kitty/graphics-protocol/#getting-the-window-size"""
+
+    buf = array.array("H", [0, 0, 0, 0])
+    fcntl.ioctl(sys.stdout, termios.TIOCGWINSZ, buf)
+    num_rows, _, _, screen_height = buf
+
+    return int(screen_height // num_rows)
+
+
+def num_required_lines(img_buf):
+    with Image.open(img_buf) as img:
+        _, img_height = img.size
+        img_buf.seek(0)
+
+    return math.ceil(img_height / get_char_cell_height())