[OSAB] Add basic scatterplot test

robert · robert · commit ada3c9c48e07 · 2019-03-31T12:54:01.000-07:00
diff --git a/bashplotlib/scatterplot.py b/bashplotlib/scatterplot.py
@@ -9,6 +9,7 @@
 import csv
 import sys
 import optparse
+import os
 from .utils.helpers import *
 from .utils.commandhelp import scatter
 
@@ -28,27 +29,27 @@ def get_scale(series, is_y=False, steps=20):
     return scaled_series
 
 
-def _plot_scatter(xs, ys, size, pch, colour, title, cs):
+def build_scatter(xs, ys, size, pch, colour, title):
     plotted = set()
 
     plot_str = ""
 
     if title:
         plot_str += box_text(title, 2 * (len(get_scale(xs, False, size)) + 1))
+        plot_str += os.linesep
 
     plot_str += "-" * (2 * (len(get_scale(xs, False, size)) + 2))
+    plot_str += os.linesep
     for y in get_scale(ys, True, size):
-        plot_str += "|"
+        plot_str += "| "
         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)
-        plot_str += " |"
+            plot_str += buildcolour(point + " ", colour)
+        plot_str += " |" + os.linesep
 
     plot_str += "-" * (2 * (len(get_scale(xs, False, size)) + 2))
     return plot_str
@@ -66,7 +67,6 @@ def plot_scatter(f, xs, ys, size, pch, colour, title):
         colour -- colour of the points
         title -- title of the plot
     """
-    cs = None
     if f:
         if isinstance(f, str):
             with open(f) as fh:
@@ -75,8 +75,6 @@ def plot_scatter(f, xs, ys, size, pch, colour, title):
             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
     else:
@@ -85,8 +83,7 @@ def plot_scatter(f, xs, ys, size, pch, colour, title):
         with open(ys) as fh:
             ys = [float(str(row).strip()) for row in fh]
 
-    print _plot_scatter(xs, ys, size, pch, colour, title, cs)
-    
+    print(build_scatter(xs, ys, size, pch, colour, title))
 
 
 def main():
diff --git a/bashplotlib/utils/helpers.py b/bashplotlib/utils/helpers.py
@@ -30,18 +30,30 @@ def get_colour(colour):
     """
     Get the escape code sequence for a colour
     """
-    return bcolours.get(colour, bcolours['ENDC'])
+    return bcolours.get(colour, bcolours['default'])
 
 
-def printcolour(text, sameline=False, colour=get_colour("ENDC")):
+def printcolour(text, sameline=False, colour=get_colour("default")):
     """
-    Print color text using escape codes
+    Print color text
     """
     if sameline:
         sep = ''
     else:
         sep = '\n'
-    sys.stdout.write(get_colour(colour) + text + bcolours["ENDC"] + sep)
+    sys.stdout.write(buildcolour(text, colour) + sep)
+
+
+def buildcolour(text, colour="default"):
+    """
+    Build color text using escape codes
+    """
+    # Unwise hack to avoid having to deal with colors in the
+    # tests.
+    if colour == "default":
+        return text
+    else:
+        return get_colour(colour) + text + bcolours["ENDC"]
 
 
 def drange(start, stop, step=1.0, include_stop=False):
diff --git a/test.py b/test.py
@@ -0,0 +1,44 @@
+from bashplotlib.scatterplot import build_scatter
+
+if __name__ == "__main__":
+    # Build a scatter plot
+    scatter = build_scatter(
+            [-10,20,30],
+            [-10,20,30],
+            10,
+            'x',
+            'default',
+            'test hello')
+
+    # This is what we know our scatter plot should
+    # look like. We have to be careful that indentations
+    # match.
+    expected_scatter = """----------------------------
+|        test hello        |
+----------------------------
+----------------------------
+|                       x  |
+|                          |
+|                          |
+|                   x      |
+|                          |
+|                          |
+|                          |
+|                          |
+|                          |
+|                          |
+|                          |
+| x                        |
+----------------------------"""
+
+    # Compare what we got to what we expected, and
+    # fail with a helpful error message if anything is
+    # different.
+    if expected_scatter == scatter:
+        print("SUCCESS!")
+    else:
+        print("FAILURE!")
+        print("Expected:")
+        print(expected_scatter)
+        print("Got:")
+        print(scatter)
