A plotter... for terminal!

termplot.py

import curses
import decimal
from math import *

log = []

def drange(start, stop, step):
    start = decimal.Decimal(str(start))
    step = decimal.Decimal(str(step))
    while start < stop:
        yield float(start)
        start += step

def addpoint(win, x, y, c='⦁', zoom=1, cp=None):
    ymax, xmax = win.getmaxyx()
    calc_y = ymax//2 - int(y*zoom)
    calc_x = xmax//2 + int(x*zoom)
    try:
        if cp:
            win.addstr(calc_y, calc_x, c, curses.color_pair(cp))
        else:
            win.addstr(calc_y, calc_x, c)
        win.noutrefresh()
    except Exception as e:
        pass

def text_corner(win, txt):
    ymax, xmax = win.getmaxyx()
    calc_x = xmax - 5 - len(txt)
    win.addstr(0, calc_x, '┌' + '─'*(len(txt) + 2) + '┐')
    win.addstr(1, calc_x, '│ ')
    win.addstr(1, calc_x + 2, txt, curses.color_pair(1))
    win.addstr(1, xmax - 3, ' │')
    win.addstr(2, calc_x, '└' + '─'*(len(txt) + 2) + '┘')
    win.noutrefresh()

def write_axes(win, xoffset=0, yoffset=0, zoom=1):
    ymax, xmax = win.getmaxyx()
    calc_y = ymax//2 - int(yoffset * zoom)
    calc_x = xmax//2 + int(xoffset * zoom)
    if calc_y < 1: calc_y = 1
    if calc_y >= (ymax - 2): calc_y = ymax - 2
    if calc_x < 0: calc_x = 0
    if calc_x >= (xmax - 2): calc_x = xmax - 2
    for x in range(xmax):
        win.addstr(calc_y, x, '─')
    for y in range(ymax):
        win.addstr(y, calc_x, '│')
    win.addstr(calc_y, calc_x, '┼')

def intro(stdscr):
    ymax, xmax = stdscr.getmaxyx()
    msg = 'Welcome to Terminal Plotter'
    msg2 = 'Enter a function below to get started:'
    stdscr.addstr(0, 0, msg)
    stdscr.addstr(1, 0, msg2)
    stdscr.addstr(3, 0, '‾'*20)
    curses.echo()
    try:
        s = stdscr.getstr(2, 0, 20)
    except KeyboardInterrupt:
        import sys
        sys.exit(0)
    curses.noecho()
    return read_fn(s)

def mouse_mode(stdscr, fn, xoffset, yoffset, zoom, ch = '+'):
    ymax, xmax = stdscr.getmaxyx()
    pos = [-xoffset, -yoffset]
    while True:
        stdscr.erase()
        write_axes(stdscr, xoffset, yoffset, zoom)
        i = 1
        for x, y in zip(drange(-xmax//2, xmax//2, 1/zoom), map(fn, drange(-xmax//2, xmax//2, 1/zoom))):
            addpoint(stdscr, x + xoffset, y + yoffset, zoom=zoom)
        addpoint(stdscr, pos[0] + xoffset, pos[1] + yoffset, zoom=zoom, c=ch, cp=2)
        text_corner(stdscr, f'Position: ({round(pos[0], 2)}, {round(pos[1], 2)})')
        curses.doupdate()
        c = stdscr.getkey()
        if c == 'q':
            break
        elif c == 'KEY_LEFT':
            pos[0] -= 1/zoom
        elif c == 'KEY_RIGHT':
            pos[0] += 1/zoom
        elif c == 'KEY_UP':
            pos[1] += 1/zoom
        elif c == 'KEY_DOWN':
            pos[1] -= 1/zoom

def plot_fn(stdscr, fn):
    ymax, xmax = stdscr.getmaxyx()
    curses.curs_set(0)
    curses.init_pair(1, curses.COLOR_RED, curses.COLOR_BLACK)
    curses.init_pair(2, curses.COLOR_GREEN, curses.COLOR_BLACK)
    xoffset = 0
    yoffset = 0
    zoom = 1
    while True:
        stdscr.erase()
        write_axes(stdscr, xoffset, yoffset, zoom)
        i = 1
        for x, y in zip(drange(-xmax//2, xmax//2, 1/zoom), map(fn, drange(-xmax//2, xmax//2, 1/zoom))):
            addpoint(stdscr, x + xoffset, y + yoffset, zoom=zoom)
        curses.doupdate()
        c = stdscr.getkey()
        if c == 'q':
            break
        elif c == 'KEY_LEFT':
            xoffset += 1/zoom
        elif c == 'KEY_RIGHT':
            xoffset -= 1/zoom
        elif c == 'KEY_UP':
            yoffset -= 1/zoom
        elif c == 'KEY_DOWN':
            yoffset += 1/zoom
        elif c == '=':
            zoom += 1
        elif c == '-':
            zoom -= 1 if zoom > 1 else 0
        elif c == 'm':
            mouse_mode(stdscr, fn, xoffset, yoffset, zoom)

def read_fn(s):
    return eval('lambda x: ' + s.decode())

def parabola(x):
    return -(x**2) + 5

def main(stdscr):
    fn = intro(stdscr)
    plot_fn(stdscr, fn)

if __name__ == '__main__':
    curses.wrapper(main)