From e46b4ede6dfc63fa09e5c30e7fcd76e8c1041044 Mon Sep 17 00:00:00 2001
From: Jarrod Johnson <jjohnson2@lenovo.com>
Date: Wed, 16 Apr 2025 12:50:59 -0400
Subject: [PATCH] Implement a CONFLUENT_IMAGE_PROTOCOL env variable

This directs CLI with image output to use a preferred protocol.

This is retroactively applied to stats.

Currently we prefer kitty, as it seems to be the most widely supported.

Though some things only support iterm, so that's an option.

And some only support sixel, but the user has to be the one to
figure out adding pysixel dependency.
---
 confluent_client/bin/nodeconsole | 44 +++++++++++++++++++++++++++++++-
 confluent_client/bin/stats       |  4 ++-
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/confluent_client/bin/nodeconsole b/confluent_client/bin/nodeconsole
index ddfbd0e3..aafdb1f7 100755
--- a/confluent_client/bin/nodeconsole
+++ b/confluent_client/bin/nodeconsole
@@ -32,6 +32,20 @@ import time
 import socket
 import re
 
+try:
+    # sixel is optional, attempt to import but stub out if unavailable
+    import io
+    import sixel
+
+    class DumbWriter(sixel.SixelWriter):
+        def restore_position(self, output):
+            return
+except ImportError:
+    class DumbWriter():
+        def draw(self, imgfile):
+            sys.stderr.write("PySixel not detected, Sixel format display not supported\n")
+
+
 confettypath = os.path.join(os.path.dirname(sys.argv[0]), 'confetty')
 argparser = optparse.OptionParser(
     usage="Usage: %prog [options] <noderange> [kill][-- [passthroughoptions]]",
@@ -72,6 +86,34 @@ argparser.add_option('-w','--windowed', action='store_true', default=False,
 
 (options, args) = argparser.parse_args()
 
+
+def draw_image(data):
+    imageformat = os.environ.get('CONFLUENT_IMAGE_PROTOCOL', 'kitty')
+    if imageformat == 'sixel':
+        sixel_draw(data)
+    elif imageformat == 'iterm':
+        iterm_draw(data)
+    else:
+        kitty_draw(data)
+
+
+def sixel_draw(data):
+    bindata = base64.b64decode(data)
+    binfile = io.BytesIO()
+    binfile.write(bindata)
+    binfile.seek(0)
+    DumbWriter().draw(binfile)
+
+def iterm_draw(data):
+    bindata = base64.b64decode(data)
+    datalen = len(bindata)
+    sys.stdout.write(
+        '\x1b]1337;File=inline=1;size={}:'.format(datalen))
+    sys.stdout.write(data.decode('utf8'))
+    sys.stdout.write('\a')
+    sys.stdout.write('\n')
+    sys.stdout.flush()
+
 def kitty_draw(data):
     while data:
         chunk, data = data[:4096], data[4096:]
@@ -126,7 +168,7 @@ if options.screenshot:
             imgdata = res['databynode'][node].get('image', {}).get('imgdata', None)
             if imgdata:
                 sys.stdout.write('{}: '.format(node))
-                kitty_draw(imgdata.encode())
+                draw_image(imgdata.encode())
                 sys.stdout.write('\n')
     sys.exit(0)
 
diff --git a/confluent_client/bin/stats b/confluent_client/bin/stats
index 94af75db..0893fadb 100755
--- a/confluent_client/bin/stats
+++ b/confluent_client/bin/stats
@@ -72,6 +72,8 @@ def plot(gui, output, plotdata, bins, fmt):
             tdata = io.BytesIO()
         plt.savefig(tdata)
     if not gui and not output:
+        if fmt == 'environment':
+            fmt = os.environ.get('CONFLUENT_IMAGE_PROTOCOL', 'kitty')
         if fmt == 'sixel':
             writer = DumbWriter()
             writer.draw(tdata)
@@ -108,7 +110,7 @@ aparser = argparse.ArgumentParser(description='Quick access to common statistics
 aparser.add_argument('-c', type=int, default=0, help='Column number to analyze (default is last column)')
 aparser.add_argument('-d', default=None, help='Value used to separate columns')
 aparser.add_argument('-x', default=False, action='store_true', help='Output histogram in graphical format')
-aparser.add_argument('-f', default='sixel', help='Format for histogram output (sixel/iterm/kitty)')
+aparser.add_argument('-f', default='environment', help='Format for histogram output (sixel/iterm/kitty)')
 aparser.add_argument('-s', default=0, help='Number of header lines to skip before processing')
 aparser.add_argument('-g', default=False, action='store_true', help='Open histogram in separate graphical window')
 aparser.add_argument('-o', default=None, help='Output histogram to the specified filename in PNG format')