refactoring

1 files changed, 90 insertions, 39 deletions

diff --git a/bin/utils/bar_chart.py b/bin/utils/bar_chart.py
index d7b44e2..db1532e 100644
--- a/bin/utils/bar_chart.py
+++ b/bin/utils/bar_chart.py
@@ -4,45 +4,54 @@
 # Distributed under the MIT License.
 
 import matplotlib.pyplot as plt
+from matplotlib.ticker import AutoLocator, FuncFormatter, MaxNLocator
 import numpy as np
 
 class BarChartBuilder:
-    _BAR_HEIGHT = 1.
+    _BAR_HEIGHT = .5
 
-    def __init__(self):
+    THICK_BAR_HEIGHT = _BAR_HEIGHT
+    THIN_BAR_HEIGHT = THICK_BAR_HEIGHT / 2
+
+    def __init__(self, labels_align_middle=True):
         self._fig, self._ax = plt.subplots()
+        self.labels_align_middle = labels_align_middle
 
     def set_title(self, title):
         self._ax.set_title(title)
 
-    def _get_bar_axis(self):
+    def _get_categories_axis(self):
         return self._ax.get_yaxis()
 
-    def _get_value_axis(self):
+    def _get_values_axis(self):
         return self._ax.get_xaxis()
 
-    def set_bar_axis_limits(self, start=None, end=None):
+    def set_categories_axis_limits(self, start=None, end=None):
         self._ax.set_ylim(bottom=start, top=end)
 
-    def set_value_axis_limits(self, start=None, end=None):
+    def set_values_axis_limits(self, start=None, end=None):
         self._ax.set_xlim(left=start, right=end)
 
-    def set_value_grid(self):
-        self._get_value_axis().grid()
+    def enable_grid_for_categories(self):
+        self._get_categories_axis().grid()
+
+    def enable_grid_for_values(self):
+        self._get_values_axis().grid()
 
-    def get_bar_labels(self):
-        return self._get_bar_axis().get_ticklabels()
+    def get_categories_labels(self):
+        return self._get_categories_axis().get_ticklabels()
 
-    def get_value_labels(self):
-        return self._get_value_axis().get_ticklabels()
+    def get_values_labels(self):
+        return self._get_values_axis().get_ticklabels()
 
     def set_value_label_formatter(self, fn):
-        from matplotlib.ticker import FuncFormatter
-        self._get_value_axis().set_major_formatter(FuncFormatter(fn))
+        self._get_values_axis().set_major_formatter(FuncFormatter(fn))
+
+    def any_values(self):
+        self._get_values_axis().set_major_locator(AutoLocator())
 
-    def set_integer_values_only(self):
-        from matplotlib.ticker import MaxNLocator
-        self._get_value_axis().set_major_locator(MaxNLocator(integer=True))
+    def only_integer_values(self):
+        self._get_values_axis().set_major_locator(MaxNLocator(integer=True))
 
     @staticmethod
     def set_property(*args, **kwargs):
@@ -60,35 +69,32 @@ class BarChartBuilder:
     def set_height(self, inches):
         self._set_size(inches, dim=1)
 
-    def plot_bars(
-            self, bar_labels, bar_lengths,
-            bars_between_ticks=False,
-            inches_per_bar=1):
+    def plot_bars(self, categories, values, inches_per_bar=THICK_BAR_HEIGHT):
 
-        numof_bars = len(bar_labels)
+        numof_bars = len(categories)
 
         if not numof_bars:
-            self.set_height(1)
-            self._get_bar_axis().set_tick_params(labelleft=False)
+            self.set_height(2 * inches_per_bar)
+            self._get_categories_axis().set_tick_params(labelleft=False)
             return []
 
-        self.set_height(inches_per_bar * numof_bars)
+        categories_axis_min = 0
+        categories_axis_max = 2 * inches_per_bar * numof_bars
 
-        bar_offsets = np.arange(numof_bars) * 2 * self._BAR_HEIGHT + self._BAR_HEIGHT
+        self.set_height(categories_axis_max)
+        self.set_categories_axis_limits(categories_axis_min, categories_axis_max)
 
-        if bars_between_ticks:
-            self._get_bar_axis().set_ticks(bar_offsets - self._BAR_HEIGHT)
-        else:
-            self._get_bar_axis().set_ticks(bar_offsets)
+        bar_offsets = 2 * inches_per_bar * np.arange(numof_bars) + inches_per_bar
 
-        bar_axis_min = 0
-        bar_axis_max = 2 * self._BAR_HEIGHT * numof_bars
-        self.set_bar_axis_limits(bar_axis_min, bar_axis_max)
+        if self.labels_align_middle:
+            self._get_categories_axis().set_ticks(bar_offsets)
+        else:
+            self._get_categories_axis().set_ticks(bar_offsets - inches_per_bar)
 
-        self._get_bar_axis().set_ticklabels(bar_labels)
+        self._get_categories_axis().set_ticklabels(categories)
 
-        return self._ax.barh(
-            bar_offsets, bar_lengths, align='center', height=self._BAR_HEIGHT)
+        return self._ax.barh(bar_offsets, values, align='center',
+                             height=inches_per_bar)
 
     @staticmethod
     def show():
@@ -98,6 +104,51 @@ class BarChartBuilder:
         self._fig.savefig(path, bbox_inches='tight')
 
 if __name__ == '__main__':
-    builder = BarChartBuilder()
-    builder.plot_bars(['Name #1', 'Name #2', 'Name #3'], [4, 11, 7])
-    builder.show()
+    import argparse
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--categories', nargs='*', default=[],
+                        help='categories')
+    parser.add_argument('--values', nargs='*', type=float, default=[],
+                        help='values')
+
+    parser.add_argument('--output', '-o',
+                        help='set output file path')
+
+    parser.add_argument('--middle', action='store_true',
+                        dest='labels_align_middle',
+                        help='align labels to the middle of the bars')
+
+    parser.add_argument('--integer-values', action='store_true',
+                        dest='only_integer_values')
+    parser.add_argument('--any-values', action='store_false',
+                        dest='only_integer_values')
+
+    parser.add_argument('--grid-categories', action='store_true')
+    parser.add_argument('--grid-values', action='store_true')
+
+    args = parser.parse_args()
+
+    if len(args.categories) < len(args.values):
+        parser.error('too many bar values')
+    if len(args.categories) > len(args.values):
+        args.values.extend([0.0] * (len(args.categories) - len(args.values)))
+
+    builder = BarChartBuilder(labels_align_middle=args.labels_align_middle)
+
+    if args.only_integer_values:
+        builder.only_integer_values()
+    else:
+        builder.any_values()
+
+    if args.grid_categories:
+        builder.enable_grid_for_categories()
+    if args.grid_values:
+        builder.enable_grid_for_values()
+
+    builder.plot_bars(args.categories, args.values)
+
+    if args.output is None:
+        builder.show()
+    else:
+        builder.save(args.output)