[canvas- graph-] draw lines at user-defined x or y

visidata/canvas.py

@@ -234,12 +234,28 @@ def rowsWithin(self, plotter_bbox, invert_y=False):
 
     def draw(self, scr):
         windowHeight, windowWidth = scr.getmaxyx()
-        disp_canvas_charset = self.options.disp_canvas_charset or ' o'
-        disp_canvas_charset += (256 - len(disp_canvas_charset)) * disp_canvas_charset[-1]
-
         if self.needsRefresh:
             self.render(windowHeight, windowWidth)
 
+        self.draw_pixels(scr)
+        self.draw_labels(scr)
+
+    def draw_empty(self, scr):
+        # use draw_empty() when calling draw_pixels() with clear_empty_squares=False
+        cursorBBox = self.plotterCursorBox
+        for char_y in range(0, self.plotheight//4):
+            for char_x in range(0, self.plotwidth//2):
+                cattr = ColorAttr()
+                ch = ' '
+                # draw cursor
+                if cursorBBox.contains(char_x*2, char_y*4) or \
+                    cursorBBox.contains(char_x*2+1, char_y*4+3):
+                    cattr = update_attr(cattr, colors.color_current_row)
+                scr.addstr(char_y, char_x, ch, cattr.attr)
+
+    def draw_pixels(self, scr, clear_empty_squares=True):
+        disp_canvas_charset = self.options.disp_canvas_charset or ' o'
+        disp_canvas_charset += (256 - len(disp_canvas_charset)) * disp_canvas_charset[-1]
         if self.pixels:
             cursorBBox = self.plotterCursorBox
             getPixelAttr = self.getPixelAttrRandom if self.options.disp_pixel_random else self.getPixelAttrMost
@@ -264,19 +280,26 @@ def draw(self, scr):
                             braille_num += pow2
                         pow2 *= 2
 
+                    ch = disp_canvas_charset[braille_num]
                     if braille_num != 0:
                         color = Counter(c for c in block_attrs if c).most_common(1)[0][0]
                         cattr = colors.get_color(color)
                     else:
                         cattr = ColorAttr()
+                        # don't erase empty squares, useful for subclasses that draw elements like gridlines
+                        # before pixels are drawn
+                        if not clear_empty_squares:
+                            continue
 
+                    # draw cursor
                     if cursorBBox.contains(char_x*2, char_y*4) or \
                        cursorBBox.contains(char_x*2+1, char_y*4+3):
                         cattr = update_attr(cattr, colors.color_current_row)
 
                     if cattr.attr:
-                        scr.addstr(char_y, char_x, disp_canvas_charset[braille_num], cattr.attr)
+                        scr.addstr(char_y, char_x, ch, cattr.attr)
 
+    def draw_labels(self, scr):
         def _mark_overlap_text(labels, textobj):
             def _overlaps(a, b):
                 a_x1, _, a_txt, _, _ = a
@@ -318,7 +341,7 @@ def _overlaps(a, b):
                         cursorBBox = self.plotterCursorBox
                         for c in txt:
                             w = dispwidth(c)
-                            # check if the cursor contains the midpoint of the character box
+                            # draw cursor if the cursor contains the midpoint of the character cell
                             if cursorBBox.contains(char_x*2+1, char_y*4+2):
                                 char_attr = update_attr(cattr, colors.color_current_row)
                                 clipdraw(scr, char_y, char_x, c, char_attr, w)
@@ -743,7 +766,6 @@ def deleteSourceRows(self, rows):
         self.source.deleteBy(lambda r,rows=rows: r in rows)
         self.reload()
 
-
 Plotter.addCommand('v', 'visibility', 'options.disp_graph_labels = not options.disp_graph_labels', 'toggle disp_graph_labels option')
 
 Canvas.addCommand(None, 'go-left', 'if cursorBox: sheet.cursorBox.xmin -= cursorBox.w', 'move cursor left by its width')

visidata/graph.py

@@ -1,10 +1,16 @@
 import math
 
 from visidata import VisiData, Canvas, Sheet, Progress, BoundingBox, Point, ColumnsSheet
-from visidata import vd, asyncthread, dispwidth, colors, clipstr
+from visidata import vd, asyncthread, dispwidth, colors, clipstr, ColorAttr, update_attr
+from visidata.type_date import date
+from statistics import median
 
 vd.theme_option('color_graph_axis', 'bold', 'color for graph axis labels')
 vd.theme_option('disp_graph_tick_x', '╵', 'character for graph x-axis ticks')
+vd.theme_option('color_gridline', '', 'color for graph grid lines')
+vd.theme_option('disp_graph_lines_x_charset', '▏││▕', 'charset to render vertical lines on graph')
+vd.theme_option('disp_graph_lines_y_charset', '▔──▁', 'charset to render horizontal lines on graph')
+vd.theme_option('disp_graph_multiple_lines_char', '▒', 'char to render multiple parallel lines')
 
 
 @VisiData.api
@@ -69,6 +75,11 @@ def __init__(self, *names, **kwargs):
         self.ylabel_maxw = 0
         super().__init__(*names, **kwargs)
 
+        self.gridlines_x = []
+        self.gridlines_y = []
+        self.gridlines_char_x = {}    # { x value in character coordinates -> character to use to draw that vertical line }
+        self.gridlines_char_y = {}    # { y value in character coordinates -> character to use to draw that horizontal line }
+
         vd.numericCols(self.xcols) or vd.fail('at least one numeric key col necessary for x-axis')
         self.ycols or vd.fail('%s is non-numeric' % '/'.join(yc.name for yc in kwargs.get('ycols')))
 
@@ -113,6 +124,39 @@ def reload(self):
         self.resetBounds()
         self.refresh()
 
+    def draw(self, scr):
+        windowHeight, windowWidth = scr.getmaxyx()
+        if self.needsRefresh:
+            self.render(windowHeight, windowWidth)
+
+        # required because we use clear_empty_squares=False for draw_pixels()
+        self.draw_empty(scr)
+        # draw gridlines first so pixels draw over them
+        self.draw_gridlines(scr)
+        # use clear_empty_squares to keep gridlines
+        self.draw_pixels(scr, clear_empty_squares=False)
+        self.draw_labels(scr)
+
+    def draw_gridlines(self, scr):
+        cursorBBox = self.plotterCursorBox
+        # draws only on character cells that have gridlines, leaves other cells unaffected
+        for char_y in range(0, self.plotheight//4):
+            has_y_line = char_y in self.gridlines_char_y.keys()
+            for char_x in range(0, self.plotwidth//2):
+                has_x_line = char_x in self.gridlines_char_x.keys()
+                if has_x_line or has_y_line:
+                    cattr = colors.color_gridline
+                    if has_x_line:
+                        ch = self.gridlines_char_x[char_x]
+                        # where two lines cross, draw the vertical line, not the horizontal one
+                    elif has_y_line:
+                        ch = self.gridlines_char_y[char_y]
+                    # draw cursor
+                    if cursorBBox.contains(char_x*2, char_y*4) or \
+                        cursorBBox.contains(char_x*2+1, char_y*4+3):
+                        cattr = update_attr(cattr, colors.color_current_row)
+                    scr.addstr(char_y, char_x, ch, cattr.attr)
+
     def resetBounds(self):
         super().resetBounds()
         self.createLabels()
@@ -123,6 +167,43 @@ def moveToRow(self, rowstr):
         self.cursorBox.h = ymax-ymin
         return True
 
+    def plot_elements(self, invert_y=True):
+        self.plot_gridlines()
+        super().plot_elements(invert_y=True)
+
+    def plot_gridlines(self):
+        self.gridlines_char_x = {}
+        self.gridlines_char_y = {}
+
+        bb = self.visibleBox
+        xmin, ymin, xmax, ymax = bb.xmin, bb.ymin, bb.xmax, bb.ymax
+
+        for data_y in self.gridlines_y:
+            data_y = float(data_y)
+            if data_y >= ymin and data_y <= ymax:
+                char_y, offset = divmod(self.scaleY(data_y), 4)
+                chars = self.options.disp_graph_lines_y_charset
+                # if we're drawing two different guide_lines in the same square, fill it with a different char
+                if char_y in self.gridlines_char_y and self.gridlines_char_y[char_y] != chars[offset]:
+                    self.gridlines_char_y[char_y] = vd.options.disp_graph_multiple_lines_char
+                else:
+                    self.gridlines_char_y[char_y] = chars[offset]
+
+        for data_x in self.gridlines_x:
+            data_x = float(data_x)
+            if data_x >= xmin and data_x <= xmax:
+                plot_x = self.scaleX(data_x)
+                # plot_x is an integer count of plotter pixels, and each character box has 2 plotter pixels
+                char_x = plot_x // 2
+                # To subdivide the 2 plotter pixels per square into 4 zones, we have to first multiply by 2.
+                offset = 2*plot_x % 4
+                chars = self.options.disp_graph_lines_x_charset
+                # if we're drawing two different guide_lines in the same square, fill it with a different char
+                if char_x in self.gridlines_char_x and self.gridlines_char_x[char_x] != chars[offset]:
+                    self.gridlines_char_x[char_x] = vd.options.disp_graph_lines_crossing
+                else:
+                    self.gridlines_char_x[char_x] = chars[offset]
+
     def moveToCol(self, colstr):
         xmin, xmax = map(float, map(self.parseX, colstr.split()))
         self.cursorBox.xmin = xmin
@@ -226,6 +307,58 @@ def rowsWithin(self, plotter_bbox):
         rows = super().rowsWithin(plotter_bbox)
         return sorted(rows, key=lambda r: self.row_order[self.source.rowid(r)])
 
+    def draw_line_x(self):
+        xcol = vd.numericCols(self.xcols)[0]
+        xtype = xcol.type
+        val = median(xcol.getValues(self.sourceRows))
+        suggested = format_input_value(val, xtype)
+        xstrs = vd.input("add line(s) at x = ", type="gridlinex", value=suggested, defaultLast=True).split()
+
+        for xstr in xstrs:
+            vals = [ v.strip() for v in xstr.split(',') ]
+            if len(vals) != len(self.xcols):
+                vd.fail(f'must have {len(self.xcols)} x values, had {len(vals)} values: {xstr}')
+            self.gridlines_x += [xtype(val) for xcol, val in zip(self.xcols, vals) if xtype(val) not in self.gridlines_x ]
+        self.refresh()
+
+    def draw_line_y(self):
+        ytype = self.ycols[0].type
+        val = median(self.ycols[0].getValues(self.sourceRows))
+        suggested = format_input_value(val, ytype)
+        ystrs = vd.input("add line(s) at y = ", type="gridliney", value=suggested, defaultLast=True).split()
+
+        self.gridlines_y += [ ytype(y) for y in ystrs if ytype(y) not in self.gridlines_y ]
+        self.refresh()
+
+    def erase_line_x(self):
+        if len(self.gridlines_x) == 0:
+            vd.fail(f'no x gridline to erase')
+        xtype = vd.numericCols(self.xcols)[0].type
+        suggested = format_input_value(self.gridlines_x[0], xtype)
+
+        xstrs = vd.input('remove line(s) at x = ', value=suggested, type='gridlinex', defaultLast=True).split()
+        for input_x in xstrs:
+            self.gridlines_x.remove(xtype(input_x))
+        self.refresh()
+
+    def erase_line_y(self):
+        ytype = self.ycols[0].type
+        suggested = format_input_value(self.gridlines_y[0], ytype) if self.gridlines_y else ''
+        ystrs = vd.input('remove line(s) at y = ', value=suggested, type='gridliney', defaultLast=True).split()
+        for y in ystrs:
+            try:
+                self.gridlines_y.remove(ytype(y))
+            except ValueError:
+                vd.fail(f'value {y} not in gridlines_y')
+        self.refresh()
+
+def format_input_value(val, type):
+    '''format a value for entry into vd.input(), so its representation has no spaces and no commas'''
+    if type is date:
+        return val.strftime('%Y-%m-%d')
+    else:
+        return str(val)
+
 
 Sheet.addCommand('.', 'plot-column', 'vd.push(GraphSheet(sheet.name, "graph", source=sheet, sourceRows=rows, xcols=keyCols, ycols=numericCols([cursorCol])))', 'plot current numeric column vs key columns; numeric key column is used for x-axis, while categorical key columns determine color')
 Sheet.addCommand('g.', 'plot-numerics', 'vd.push(GraphSheet(sheet.name, "graph", source=sheet, sourceRows=rows, xcols=keyCols, ycols=numericCols(nonKeyVisibleCols)))', 'plot a graph of all visible numeric columns vs key columns')
@@ -261,6 +394,12 @@ def set_x(sheet, s):
 Canvas.addCommand('y', 'resize-y-input', 'sheet.set_y(input("set ymin ymax="))', 'set ymin/ymax on graph axes')
 Canvas.addCommand('x', 'resize-x-input', 'sheet.set_x(input("set xmin xmax="))', 'set xmin/xmax on graph axes')
 
+GraphSheet.addCommand('gx', 'draw-line-x', 'sheet.draw_line_x()', 'draw a vertical line at x-values (space-separated)')
+GraphSheet.addCommand('gy', 'draw-line-y', 'sheet.draw_line_y()', 'draw a horizontal line at y-values (space-separated)')
+GraphSheet.addCommand('zx', 'erase-line-x', 'sheet.erase_line_x()', 'remove a horizontal line at x-values (space-separated)')
+GraphSheet.addCommand('zy', 'erase-line-y', 'sheet.erase_line_y()', 'remove a vertical line at y-values (space-separated)')
+GraphSheet.addCommand('gzx', 'erase-lines-x', 'sheet.gridlines_x = []; sheet.refresh()', 'erase all vertical x-value lines')
+GraphSheet.addCommand('gzy', 'erase-lines-y', 'sheet.gridlines_y = []; sheet.refresh()', 'erase any horizontal y-value lines')
 
 vd.addGlobals({
     'GraphSheet': GraphSheet,

visidata/themes/ascii8.py

@@ -81,5 +81,7 @@
     disp_menu_input='_',
     disp_menu_fmt='7-bit ASCII 3-bit color',
     plot_colors = 'white',
-    disp_histogram='*'
+    disp_histogram='*',
+    disp_graph_lines_x_charset='||||',
+    disp_graph_lines_y_charset='----'
 )