bug fix for scatterplot: 1). add delimiter, default to tab; 2). default scatter colour to None to remove strange characters in pipe output; 3). fix broken pipe during piping

Michael Wang · Michael Wang · commit 767c14f093fc · 2020-08-22T10:42:03.000-07:00
diff --git a/bashplotlib/scatterplot.py b/bashplotlib/scatterplot.py
@@ -12,13 +12,15 @@
 from .utils.helpers import *
 from .utils.commandhelp import scatter
 
+from signal import signal, SIGPIPE, SIG_DFL
+signal(SIGPIPE,SIG_DFL)
+
 
 def get_scale(series, is_y=False, steps=20):
     min_val = min(series)
     max_val = max(series)
     scaled_series = []
-    for x in drange(min_val, max_val, (max_val - min_val) / steps,
-                    include_stop=True):
+    for x in drange(min_val, max_val, (max_val - min_val) / steps):
         if x > 0 and scaled_series and max(scaled_series) < 0:
             scaled_series.append(0.0)
         scaled_series.append(x)
@@ -28,61 +30,58 @@ def get_scale(series, is_y=False, steps=20):
     return scaled_series
 
 
-def _plot_scatter(xs, ys, size, pch, colour, title, cs):
-    plotted = set()
-
-    if title:
-        print(box_text(title, 2 * len(get_scale(xs, False, size)) + 1))
-
-    print("-" * (2 * len(get_scale(xs, False, size)) + 2))
-    for y in get_scale(ys, True, size):
-        print("|", end=' ')
-        for x in get_scale(xs, False, size):
-            point = " "
-            for (i, (xp, yp)) in enumerate(zip(xs, ys)):
-                if xp <= x and yp >= y and (xp, yp) not in plotted:
-                    point = pch
-                    plotted.add((xp, yp))
-                    if cs:
-                        colour = cs[i]
-            printcolour(point, True, colour)
-        print(" |")
-    print("-" * (2 * len(get_scale(xs, False, size)) + 2))
-
-def plot_scatter(f, xs, ys, size, pch, colour, title):
+def plot_scatter(f, xs, ys, size, pch, colour, title, delimiter):
     """
     Form a complex number.
 
     Arguments:
-        f -- comma delimited file w/ x,y coordinates
+        f -- delimited (by delimiter) file w/ x,y coordinates
         xs -- if f not specified this is a file w/ x coordinates
         ys -- if f not specified this is a filew / y coordinates
         size -- size of the plot
         pch -- shape of the points (any character)
         colour -- colour of the points
         title -- title of the plot
     """
-    cs = None
+
     if f:
         if isinstance(f, str):
-            with open(f) as fh:
-                data = [tuple(line.strip().split(',')) for line in fh]
-        else:
-            data = [tuple(line.strip().split(',')) for line in f]
-        xs = [float(i[0]) for i in data]
-        ys = [float(i[1]) for i in data]
-        if len(data[0]) > 2:
-            cs = [i[2].strip() for i in data]
-    elif isinstance(xs, list) and isinstance(ys, list):
-        pass
+            f = open(f)
+
+        data = [tuple(map(float, line.strip().split(delimiter))) for line in f]
+        xs = [i[0] for i in data]
+        ys = [i[1] for i in data]
     else:
-        with open(xs) as fh:
-            xs = [float(str(row).strip()) for row in fh]
-        with open(ys) as fh:
-            ys = [float(str(row).strip()) for row in fh]
+        xs = [float(str(row).strip()) for row in open(xs)]
+        ys = [float(str(row).strip()) for row in open(ys)]
 
-    _plot_scatter(xs, ys, size, pch, colour, title, cs)
-    
+    plotted = set()
+
+    if title:
+        print(box_text(title, 2 * len(get_scale(xs, False, size)) + 1))
+
+    print("-" * (2 * len(get_scale(xs, False, size)) + 2))
+    for y in get_scale(ys, True, size):
+        print("|", end=' ')
+        for x in get_scale(xs, False, size):
+            point = " "
+            for (i, (xp, yp)) in enumerate(zip(xs, ys)):
+                if xp <= x and yp >= y and (xp, yp) not in plotted:
+                    point = pch
+                    #point = str(i)
+                    plotted.add((xp, yp))
+            if x == 0 and y == 0:
+                point = "o"
+            elif x == 0:
+                point = "|"
+            elif y == 0:
+                point = "-"
+			if colour:
+            	printcolour(point, True, colour)
+            else:
+				print(point, end='')
+        print("|")
+    print("-" * (2 * len(get_scale(xs, False, size)) + 2))
 
 
 def main():
@@ -94,17 +93,18 @@ def main():
     parser.add_option('-x', help='x coordinates', default=None, dest='x')
     parser.add_option('-y', help='y coordinates', default=None, dest='y')
     parser.add_option('-s', '--size', help='y coordinates', default=20, dest='size', type='int')
+    parser.add_option('-d', '--delimiter', help='delimiter', default="\t", dest='delimiter')
     parser.add_option('-p', '--pch', help='shape of point', default="x", dest='pch')
     parser.add_option('-c', '--colour', help='colour of the plot (%s)' %
-                      colour_help, default='default', dest='colour')
+                      colour_help, default=None, dest='colour')
 
     opts, args = parser.parse_args()
 
     if opts.f is None and (opts.x is None or opts.y is None):
         opts.f = sys.stdin.readlines()
 
     if opts.f or (opts.x and opts.y):
-        plot_scatter(opts.f, opts.x, opts.y, opts.size, opts.pch, opts.colour, opts.t)
+        plot_scatter(opts.f, opts.x, opts.y, opts.size, opts.pch, opts.colour, opts.t, opts.delimiter)
     else:
         print("nothing to plot!")
 
diff --git a/bashplotlib/utils/helpers.py b/bashplotlib/utils/helpers.py
@@ -20,30 +20,29 @@
     "grey":    '\033[90m',
     "black":   '\033[30m',
     "default": '\033[39m',
+	"none":    '',
     "ENDC":    '\033[39m',
 }
 
 colour_help = ', '.join([colour for colour in bcolours if colour != "ENDC"])
 
 
-def get_colour(colour, default="default"):
+def get_colour(colour):
     """
     Get the escape code sequence for a colour
     """
-    return bcolours.get(colour, bcolours[default])
+    return bcolours.get(colour, bcolours['ENDC'])
 
 
-def printcolour(text, sameline=False, colour="default"):
+def printcolour(text, sameline=False, colour=get_colour("ENDC")):
     """
     Print color text using escape codes
     """
-    sep = '' if sameline else '\n'
-
-    # If no colour set, do not print color ESC characters
-    if get_colour(colour) == get_colour("ENDC"):
-        sys.stdout.write(text + sep)
+    if sameline:
+        sep = ''
     else:
-        sys.stdout.write(get_colour(colour) + text + get_colour("ENDC") + sep)
+        sep = '\n'
+    sys.stdout.write(get_colour(colour) + text + bcolours["ENDC"] + sep)
 
 
 def drange(start, stop, step=1.0, include_stop=False):
@@ -66,18 +65,6 @@ def drange(start, stop, step=1.0, include_stop=False):
             r = round(r, 10)
 
 
-def abbreviate(labels, rfill=' '):
-    """
-    Abbreviate labels without introducing ambiguities.
-    """
-    max_len = max(len(l) for l in labels)
-    for i in range(1, max_len):
-        abbrev = [l[:i].ljust(i, rfill) for l in labels]
-        if len(abbrev) == len(set(abbrev)):
-            break
-    return abbrev
-
-
 def box_text(text, width, offset=0):
     """
     Return text inside an ascii textbox
