from contextlib import contextmanager
from csv import DictWriter
import logging

from ktbs_bench.utils.decorators import bench as util_bench


class BenchManager(object):
    """
    Manage benchmarks by timing function against different contexts.

    General concept
    ===============
    A BenchManager is instantiated in order to collect functions
    to benchmark, like so:
    >>> my_bench_manager = BenchManager()

    In order to add functions to bench, one flag them for bench
    by using the :meth:`bench` decorator. For example:
    >>> @my_bench_manager.bench
    ... def add_one(n):
    ...     return n + 1

    Each flagged function is then called against contexts.
    A context is a function with optional setup and teardown, and it
    must *yield* the parameter that benchmarked functions need.
    >>> @my_bench_manager.context
    ... def three():
    ...     # optional setup
    ...     try:
    ...        # yield the parameter
    ...        yield 3
    ...     finally:
    ...        # optional teardown
    ...        pass

    Finally, to perform the benchmarks, one must call:
    >>> my_bench_manager.run('/tmp/my_results.csv')

    The result of the two examples above is to time ``add_one(3)``.

    Technical details
    =================
    Each context is stored in the list :attr:`_contexts`.
    Each function to benchmark is stored in the list :attr:`_bench_funcs`.

    When :meth:`run` is called, it will iterate over functions and contexts
    to call each function against each context.
    """

    def __init__(self):
        self._contexts = []
        self._bench_funcs = []
        self._results = {}

    def bench(self, func):
        """Prepare a function to be benched and add it to the list to be run later.

        :param function func: the function to bench
        """
        func = util_bench(func)
        self._bench_funcs.append(func)

    def context(self, func):
        """Decorate a function to act as a context.

        :param function func: the function that describes the context
        """
        func = contextmanager(func)
        self._contexts.append(func)

    def run(self, output_filename, show_log=False):
        """Execute each collected function against each context.

        :param str output_filename: filename of the CSV output
        :param bool show_log: True to display log during the run, False otherwise
        """
        # Run the bench for each function against each context
        if show_log:
            logging.getLogger().setLevel(logging.INFO)
        for func in self._bench_funcs:
            self._results[func.__name__] = {}
            if show_log:
                logging.info('Func: %s' % func.__name__)
            for context in self._contexts:
                if show_log:
                    logging.info('with context: %s' % context.__name__)
                with context() as arg:
                    if not isinstance(arg, tuple):
                        arg = tuple([arg])  # Make arg a tuple of one element
                    _, res_time = func(*arg)
                if show_log:
                    logging.info('res time: %s' % res_time)
                self._results[func.__name__][context.__name__] = res_time

        # Write output to a CSV file
        fieldnames = ['func_name']
        fieldnames += [context.__name__ for context in self._contexts]
        with open(output_filename, 'wb') as output_file:
            csv_file = DictWriter(output_file, fieldnames)
            csv_file.writeheader()

            for func_name, env_results in self._results.items():
                row_results = env_results
                row_results['func_name'] = func_name
                csv_file.writerow(row_results)

@@ -2175,16 +2175,21 @@
 show_log
+_info
 =False):
@@ -2204,30 +2204,17 @@
  %22%22%22
-Execute each collected
+Benchmark
  fun
@@ -2210,32 +2210,33 @@
 nchmark function
+s
  against each co
@@ -2228,28 +2228,24 @@
 against 
-each 
 context
+s
 .%0A%0A     
@@ -2330,16 +2330,21 @@
 show_log
+_info
 : True t
@@ -2405,71 +2405,8 @@
 %22%22%22%0A
-        # Run the bench for each function against each context%0A
     
@@ -2412,32 +2412,37 @@
      if show_log
+_info
 :%0A            lo
@@ -2481,16 +2481,80 @@
 ng.INFO)
+%0A%0A        # Run the bench for each function against each context
 %0A       
@@ -2647,250 +2647,103 @@
     
-if show_log:%0A                logging.info('Func: %25s' %25 func.__name__)%0A            for context in self._contexts:%0A                if show_log:%0A                    logging.info('with context: %25s' %25 context.__name__)%0A                wi
+for context in self._contexts:%0A                with context() as arg:  # catch what 
 th
+e
  context
 () a
@@ -2738,26 +2738,22 @@
  context
-() as arg:
+ yield
 %0A       
@@ -2842,17 +2842,17 @@
 rg%5D)  # 
-M
+m
 ake arg 
@@ -2922,41 +2922,52 @@
 arg)
-%0A
   
-              if show_log:%0A   
+# call func with what the context yielded%0A
     
@@ -2970,33 +2970,32 @@
                 
- 
 logging.info('re
@@ -2996,23 +2996,120 @@
 fo('
-res time
+Benchmarking function %25s with context %25s
 : %25s
+ s
 ' %25
+%0A                             (func.__name__, context.__name__,
  res
@@ -3114,16 +3114,17 @@
 es_time)
+)
 %0A       
@@ -3228,16 +3228,285 @@
 SV file%0A
+        self.write_output(output_filename)%0A%0A    def write_output(self, output_filename):%0A        %22%22%22Write results of the BenchManager to a nicely formatted CSV file.%0A%0A        :param str output_filename: filename of the CSV output%0A        %22%22%22%0A        # Get column names%0A
         
@@ -3602,16 +3602,66 @@
 ontexts%5D
+%0A%0A        # Create the CSV with header and results
 %0A       

Improve code readability of BenchManager run()