1 files changed, 18 insertions, 13 deletions
diff --git a/wqflask/utility/benchmark.py b/wqflask/utility/benchmark.py
index d5b32703..8c97370d 100755..100644
--- a/wqflask/utility/benchmark.py
+++ b/wqflask/utility/benchmark.py
@@ -3,34 +3,39 @@ from __future__ import print_function, division, absolute_import
 import collections
 import inspect
 import time
+from utility.tools import LOG_BENCH
 
+from utility.logger import getLogger
+logger = getLogger(__name__ )
 
 class Bench(object):
     entries = collections.OrderedDict()
 
-    def __init__(self, name=None):
+    def __init__(self, name=None, write_output=True):
         self.name = name
+        self.write_output = write_output
 
     def __enter__(self):
-        if self.name:
-            print("Starting benchmark: %s" % (self.name))
-        else:
-            print("Starting benchmark at: %s [%i]" % (inspect.stack()[1][3], inspect.stack()[1][2]))
+        if self.write_output:
+            if self.name:
+                logger.debug("Starting benchmark: %s" % (self.name))
+            else:
+                logger.debug("Starting benchmark at: %s [%i]" % (inspect.stack()[1][3], inspect.stack()[1][2]))
         self.start_time = time.time()
 
     def __exit__(self, type, value, traceback):
-        if self.name:
-            name = self.name
-        else:
-            name = "That"
-
         time_taken = time.time() - self.start_time
-        print("  %s took: %f seconds" % (name, (time_taken)))
+        if self.write_output:
+            if self.name:
+                name = self.name
+            else:
+                name = "That"
+
+            logger.info("  %s took: %f seconds" % (name, (time_taken)))
 
         if self.name:
             Bench.entries[self.name] = Bench.entries.get(self.name, 0) + time_taken
 
-
     @classmethod
     def report(cls):
         total_time = sum((time_taken for time_taken in cls.entries.itervalues()))
@@ -42,4 +47,4 @@ class Bench(object):
 
     def reset(cls):
         """Reset the entries"""
-        cls.entries = collections.OrderedDict()
\ No newline at end of file
+        cls.entries = collections.OrderedDict()