added functionality to pass in lists

Author: kennethgoodman

bashplotlib/scatterplot.py

@@ -28,6 +28,27 @@ 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):
     """
     Form a complex number.
@@ -41,7 +62,7 @@ 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):
             f = open(f)
@@ -51,31 +72,14 @@ def plot_scatter(f, xs, ys, size, pch, colour, title):
         ys = [float(i[1]) for i in data]
         if len(data[0]) > 2:
             cs = [i[2].strip() for i in data]
-        else:
-            cs = None
+    elif isinstance(xs, list) and isinstance(ys, list):
+        pass
     else:
         xs = [float(str(row).strip()) for row in open(xs)]
         ys = [float(str(row).strip()) for row in open(ys)]
 
-    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))
+    _plot_scatter(xs, ys, size, pch, colour, title, cs)
+    
 
 
 def main():