Fix the axes, fix error with constant functions and some cleanup.

1 files changed, 22 insertions, 12 deletions

diff --git a/plotlib.py b/plotlib.py
index 7073f05..c40067d 100644
--- a/plotlib.py
+++ b/plotlib.py
@@ -199,11 +199,12 @@ class CustomPlot(_PlotBase):
 
     def draw_axes(self, labelx, labely, val):
         """Draw axes on the plot."""
-        F = 0.77
-        NOL = 5 # no of labels, atleast is 2
-        interval = (len(val) - 2)/NOL
+        F = 0.8
+        NOL = 4 # maximum no of labels
+
         y_coords = sorted([i[1] for i in val])
         x_coords = sorted([i[0] for i in val])
+
         max_y = max(y_coords)
         min_y = min(y_coords)
 
@@ -211,27 +212,36 @@ class CustomPlot(_PlotBase):
         min_x = min(x_coords)
 
         # X axis
+        interval = len(val)/(NOL - 1)
         self.plot_line((0.11, 0.89), (0.92, 0.89), "black")
         self.add_text((0.11 + min_x + F * 0, 0.93), format_float(min_x))
         plot_index = interval
 
-        while plot_index < len(val) - interval:
-            self.add_text((0.11 + F * abs(x_coords[plot_index] - min_x) / abs(max_x - min_x), 0.93), format_float(x_coords[plot_index]))
+        while plot_index <= len(val) - interval:
+            self.add_text((0.11 + F * abs(x_coords[plot_index] - min_x) / \
+                           abs(max_x - min_x), 0.93), \
+                          format_float(x_coords[plot_index]))
             plot_index += interval
 
         self.add_text((0.11 + F * 1, 0.93), format_float(max_x))
         self.add_text((0.50, 0.98), labelx)
 
         # Y axis
+        interval = float(max_y - min_y)/(NOL - 1)
         self.plot_line((0.11, 0.08), (0.11, 0.89), "black")
-        self.add_text((-0.90, 0.10), format_float(min_y), rotate=-90)
-        plot_index = interval
-
-        while plot_index < len(val) - interval:
-            self.add_text((-(0.91 - F * abs(y_coords[plot_index] - min_y) / abs(max_y - min_y)), 0.10), format_float(y_coords[plot_index]), rotate=-90)
-            plot_index += interval
+        # if its a constant function we only need to plot one label
+        if min_y == max_y:        
+            self.add_text((-0.50, 0.10), format_float(min_y), rotate=-90)        
+        else:
+            self.add_text((-0.90, 0.10), format_float(min_y), rotate=-90)        
+            plot_value = min_y + interval
+            while plot_value <= max_y - interval:
+                self.add_text((-(0.91 - F * abs(plot_value - min_y) / \
+                               abs(max_y - min_y)), 0.10), \
+                              format_float(plot_value), rotate=-90)
+                plot_value += interval
+            self.add_text((-(0.89 - F), 0.10), format_float(max_y), rotate=-90)
 
-        self.add_text((-(0.89 - F), 0.10), format_float(max_y), rotate=-90)
         self.add_text((-0.50, 0.045), labely, rotate=-90)
 
     def produce_plot(self, vals, *args, **kwargs):