Draw psql output as iTerm2 v3 inline graph using matplotlib

pg_graph

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Draw psql output as iTerm2 v3 inline graph using matplotlib
# Author: Alexander Korotkov <aekorotkov@gmail.com>
# with a few edits by nik@postgres.ai to support Python3

import sys
import re
import warnings
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import itertools
import numpy as np
from io import BytesIO
import base64
import os

# TODO:
# * handle newlines in data
# * support more output format variations: line styles, spaces etc.

# Couple of escape sequences printing functions
def print_osc():
	if os.environ.get('TERM') == 'screen':
		sys.stdout.write("\033Ptmux;\033\033]")
	else:
		sys.stdout.write("\033]")
def print_st():
	if os.environ.get('TERM') == 'screen':
		sys.stdout.write("\a\033\\")
	else:
		sys.stdout.write("\a")

enc = 'utf-8'
# Pattern for record header in extended format
extRecPattern = re.compile(u'^([─-]\[ RECORD \d+ \](?:-*|─*))(?:\+-*|[┬┼]─*)$')
data = []
colors = itertools.cycle(['b', 'g', 'r', 'c', 'm', 'y', 'k'])

# Read table from psql
header = sys.stdin.readline()
m = extRecPattern.match(header)
if m:
	# Extended format
	nameLen = len(m.group(1)) - 1
	valueOffset = len(m.group(1)) + 2
	i = 0
	colNames = []
	row = []
	line = sys.stdin.readline()
	while line:
		if not extRecPattern.match(line):
			name = line[:nameLen].rstrip(' ')
			value = line[valueOffset:].rstrip('\n')
			if i == 0:
				colNames.append(name)
			row.append(value)
		else:
			i = i + 1
			data.append(row)
			row = []
		line = sys.stdin.readline()
	data.append(row)
else:
	# Basic format
	colNames = [name.strip(' ') for name in re.split(u'[|│]', header.rstrip('\n'))]
	line = sys.stdin.readline()
	while line:
		parts = re.split(u'[|│]', line.rstrip('\n'))
		if len(parts) == len(colNames):
			row = [value.strip(' ') for value in parts]
			data.append(row)
		line = sys.stdin.readline()

# Setup graph colored white on black
fig = plt.figure(figsize = (8, 6))
ax = fig.add_subplot(1, 1, 1, facecolor='k')
ax.title.set_color('white')
ax.yaxis.label.set_color('white')
ax.xaxis.label.set_color('white')
ax.tick_params(axis = 'x', colors = 'white')
ax.tick_params(axis = 'y', colors = 'white')

# Autodetect plot type. Use plot if all the labels are floats.
x = []
try:
	for row in data:
		x.append(float(row[0]))
	chartType = 'plot'
except ValueError:
	chartType = 'bar'

if chartType == 'bar':
	for row in data:
		x.append(row[0])
	index = np.arange(len(data))
	barWidth = 1.0 / float(len(colNames))

# Do draw
for i in range(1, len(colNames)):
	color = next(colors)
	y = []
	for row in data:
		y.append(float(row[i]))
	if chartType == 'plot':
		ax.plot(x, y, label=colNames[i], color = color)
	if chartType == 'bar':
		ax.bar(index + barWidth * (float(i) - 0.5),
				y,
				barWidth,
				label = colNames[i],
				color = color)

if chartType == 'bar':
	plt.xticks(index + 0.5, x)

# Setup legend
legend = ax.legend(loc = 'best', fancybox = True, framealpha = 0.5)
legendFrame = legend.get_frame()
legendFrame.set_color('white')
legendFrame.set_facecolor('black')
legendFrame.set_edgecolor('white')
for text in legend.get_texts():
	text.set_color('white')

plt.xlabel(colNames[0])
ax.grid(True, color = 'w')

for spine in ax.spines:
	ax.spines[spine].set_edgecolor('white')

# Supress warnings in tight_layout 
with warnings.catch_warnings():
	warnings.simplefilter("ignore")
	plt.tight_layout()

# Put inlined image into iTerm2 v3
imgdata = BytesIO()
plt.savefig(imgdata, format='png', transparent=True)
size = imgdata.tell()
imgdata.seek(0)  # rewind the data

print_osc()
sys.stdout.write('1337;File=')
sys.stdout.write('name=' + base64.b64encode('graph.png') + ';')
sys.stdout.write('size=' + str('size') + ';')
sys.stdout.write('inline=1:')
sys.stdout.write(base64.b64encode(imgdata.buf))
print_st()
sys.stdout.write("\n")

pg_graph_quickchart_io.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Draw chart as iTerm2 v3 inline graph using QuickChart.io
# Original idea by Alexander Korotkov: https://akorotkov.github.io/blog/2016/06/09/psql-graph/
# Adapted for QuickChart.io

# See: https://twitter.com/samokhvalov/status/1736301721155191218

import sys
import os
import json
import base64
import requests

def print_osc():
    if os.environ.get('TERM') == 'screen':
        sys.stdout.write("\033Ptmux;\033\033]")
    else:
        sys.stdout.write("\033]")

def print_st():
    if os.environ.get('TERM') == 'screen':
        sys.stdout.write("\a\033\\")
    else:
        sys.stdout.write("\a")

# Read and parse input data
headers = []
data = []
for line in sys.stdin:
    line = line.strip()
    if not line or '-' in line or line.startswith('('):  # Skip empty lines, the separator line, and summary lines
        continue
    row = [item.strip() for item in line.split('|')]
    if not headers:
        headers = row
    else:
        data.append(row)

# Validate data
if not headers or not data:
    print("No data found", file=sys.stderr)
    sys.exit(1)

# Check each data row for correct number of elements
for row in data:
    if len(row) != len(headers):
        print(f"Data row length does not match header length. Row: {row}", file=sys.stderr)
        sys.exit(1)

# Construct chart configuration for QuickChart.io
chart_config = {
    'type': 'bar',  # Change to 'line' or other types as needed
    'data': {
        'labels': [row[0] for row in data],
        'datasets': [{'label': headers[i], 'data': [float(row[i]) for row in data]} for i in range(1, len(headers))]
    }
}

url = f"https://quickchart.io/chart?c={json.dumps(chart_config)}"

# Fetch the chart image
response = requests.get(url)
if response.status_code != 200:
    print("Failed to retrieve image", file=sys.stderr)
    sys.exit(1)

# Output the image as an inline image in iTerm2
print_osc()
sys.stdout.write('1337;File=')
sys.stdout.write('name=' + base64.b64encode('chart.png'.encode()).decode() + ';')
sys.stdout.write('size=' + str(len(response.content)) + ';')
sys.stdout.write('inline=1:')
sys.stdout.write(base64.b64encode(response.content).decode())
print_st()
sys.stdout.write("\n")