diff --git a/bashplotlib/scatterplot.py b/bashplotlib/scatterplot.py index 69cab9d..bcaa3c8 100644 --- a/bashplotlib/scatterplot.py +++ b/bashplotlib/scatterplot.py @@ -14,8 +14,11 @@ def get_scale(series, is_y=False, steps=20): - min_val = min(series) - max_val = max(series) + min_val = -1 * (steps / 2) + max_val = steps / 2 + if series != []: + 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): @@ -28,15 +31,20 @@ def get_scale(series, is_y=False, steps=20): return scaled_series -def _plot_scatter(xs, ys, size, pch, colour, title, cs): +def _plot_scatter(xs, ys, size, pch, title, cs, x_title, y_title): plotted = set() + graph = "" if title: - print(box_text(title, 2 * (len(get_scale(xs, False, size)) + 1))) + graph += box_text(title, 2 * (len(get_scale(xs, False, size)) + 1)) + '\n' - print("-" * (2 * (len(get_scale(xs, False, size)) + 2))) + graph += " +" + "-" * (2 * (len(get_scale(xs, False, size)) + 2) - 2) + "+\n" for y in get_scale(ys, True, size): - print("|", end=' ') + if y_title == "": + graph += " | " + else: + graph += y_title[:1] + " | " + y_title = y_title[1:] for x in get_scale(xs, False, size): point = " " for (i, (xp, yp)) in enumerate(zip(xs, ys)): @@ -45,11 +53,21 @@ def _plot_scatter(xs, ys, size, pch, colour, title, cs): 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): + # print axes + if point != pch: + if x == 0.0 and y == 0.0: + point = "o" + elif x == 0.0: + point = "|" + elif y == 0.0: + point = "-" + graph += point + " " + graph += " |\n" + graph += " +" + "-" * (2 * (len(get_scale(xs, False, size)) + 2) - 2) + "+\n" + graph += " " + x_title + return graph + +def plot_scatter(f, xs, ys, size, pch, colour, title, x_title, y_title): """ Form a complex number. @@ -81,7 +99,8 @@ def plot_scatter(f, xs, ys, size, pch, colour, title): with open(ys) as fh: ys = [float(str(row).strip()) for row in fh] - _plot_scatter(xs, ys, size, pch, colour, title, cs) + graph = _plot_scatter(xs, ys, size, pch, title, cs, x_title, y_title) + print(graph) @@ -97,6 +116,8 @@ def main(): 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') + parser.add_option('-xt', '--xtitle', help='title for the x-axis', default="y axis", dest="x_title") + parser.add_option('-yt', '--ytitle', help='title for the y-axis', default="x axis", dest="y_title") opts, args = parser.parse_args() @@ -104,7 +125,7 @@ def main(): 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.x_title, opts.y_title) else: print("nothing to plot!") diff --git a/bashplotlib/utils/helpers.py b/bashplotlib/utils/helpers.py index cf209ee..cd37a81 100644 --- a/bashplotlib/utils/helpers.py +++ b/bashplotlib/utils/helpers.py @@ -80,7 +80,7 @@ def box_text(text, width, offset=0): """ Return text inside an ascii textbox """ - box = " " * offset + "-" * (width+2) + "\n" - box += " " * offset + "|" + text.center(width) + "|" + "\n" - box += " " * offset + "-" * (width+2) + box = " " * offset + " +" + "-" * (width) + "+\n" + box += " " * offset + " |" + text.center(width) + "|" + "\n" + box += " " * offset + " +" + "-" * (width) + "+" return box diff --git a/scratch.py b/scratch.py new file mode 100644 index 0000000..70a62d0 --- /dev/null +++ b/scratch.py @@ -0,0 +1,22 @@ +# scratch.py +from bashplotlib.scatterplot import plot_scatter + +x_coords = [-10,20,30] +y_coords = [-10,20,30] +width = 10 +char = 'x' +color = 'default' +title = 'My Test Graph' +x_title = 'My X Axis' +y_title = 'My Y Axis' + +plot_scatter( + None, + x_coords, + y_coords, + width, + char, + color, + title, + x_title, + y_title) diff --git a/test.py b/test.py new file mode 100644 index 0000000..6424074 --- /dev/null +++ b/test.py @@ -0,0 +1,69 @@ +# test.py +from bashplotlib.scatterplot import _plot_scatter +import unittest + +class TestScatterPlot(unittest.TestCase): + def test_no_points(self): + self.assertEqual(_plot_scatter([], [], 10, "x", + "My Test Graph", "default", "My X Axis", "My Y Axis"), + """ +------------------------+\n | My Test Graph | + +------------------------+ + +------------------------+ +M | | | +y | | | + | | | +Y | | | + | | | +A | - - - - - o - - - - - | +x | | | +i | | | +s | | | + | | | + | | | + +------------------------+ + My X Axis""") + + def test_points(self): + self.assertEqual(_plot_scatter([-10, -3, 20,30], [-10, 6, 20,-2], 10, "*", + "Points!", "default", "Something x", "Another y"), + """ +--------------------------+\n | Points! | + +--------------------------+ + +--------------------------+ +A | | * | +n | | | +o | | | +t | | | +h | | | +e | * | | +r | | | + | - - - o - - - - - - - - | +y | | | + | | * | + | | | + | * | | + +--------------------------+ + Something x""") + + def test_points_on_axes(self): + self.assertEqual(_plot_scatter([0, -10, 0, 20,0], [0, 0, 6, 0,-2], 10, "#", + "Points!", "default", "x", "y"), + """ +--------------------------+\n | Points! | + +--------------------------+ + +--------------------------+ +y | # | + | | | + | | | + | | | + | | | + | | | + | | | + | | | + | # - - - # - - - - - - # | + | | | + | | | + | # | + +--------------------------+ + x""") + +if __name__ == "__main__": + unittest.main()