From 01dd0fd927c299331c664662a0aac9ba21ae0fb6 Mon Sep 17 00:00:00 2001
From: Jon Dugan <jdugan@x1024.net>
Date: Fri, 2 Sep 2016 18:00:17 -0500
Subject: [PATCH] Fix decorator so that fmt works when without logging.

Make the effect of the fmt argument to the decorator consistent with
or without logging.

Added a test for this as well.
---
 contexttimer/__init__.py | 15 +++++++--------
 tests/test_timer.py      | 18 ++++++++++++++++++
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/contexttimer/__init__.py b/contexttimer/__init__.py
index 67ce2c9..d7190bc 100644
--- a/contexttimer/__init__.py
+++ b/contexttimer/__init__.py
@@ -121,18 +121,17 @@ def wrapped_f(f):
         def wrapped(*args, **kwargs):
             with Timer(**timer_kwargs) as t:
                 out = f(*args, **kwargs)
+            context = {
+                'function_name': f.__name__,
+                'execution_time': t.elapsed,
+            }
             if logger:
-                extra = {
-                    'function_name': f.__name__,
-                    'execution_time': t.elapsed,
-                }
                 logger.log(
                     level,
-                    fmt % extra,
-                    extra=extra)
+                    fmt % context,
+                    extra=context)
             else:
-                print("function %s execution time: %.3f " %
-                      (f.__name__, t.elapsed))
+                print(fmt % context)
             return out
         return wrapped
     if (len(func_or_func_args) == 1
diff --git a/tests/test_timer.py b/tests/test_timer.py
index fd0fd4e..1ea1334 100644
--- a/tests/test_timer.py
+++ b/tests/test_timer.py
@@ -27,3 +27,21 @@ def print_reversed(string):
 
             self.assertIsNotNone(output)
             self.assertRegexpMatches(output.getvalue(), expected)
+
+    def test_decorator_print(self):
+        tests = [
+            ({}, r"function foo execution time: [0-9.]+"),
+            ({'fmt': '%(execution_time)s seconds later...'}, r"[0-9.]+ seconds later..."),
+        ]
+
+
+        for kwargs, expected in tests:
+            output = StringIO()
+            with mock.patch('sys.stdout', new=output):
+                @contexttimer.timer(**kwargs)
+                def foo():
+                    pass
+                foo()
+
+            self.assertIsNotNone(output)
+            self.assertRegexpMatches(output.getvalue(), expected)