Merge pull request #98 from leplatrem/master

Bug: Fix timer decorator with partial functions (#85)

4 files changed, 28 insertions, 2 deletions

diff --git a/AUTHORS b/AUTHORS
index a37070b..bedfa5f 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -6,3 +6,4 @@ Contributors:
 - Andy McKay <amckay@mozilla.com>
 - Daniel Holz <dgholz@gmail.com>
 - Kyle Conroy <kyle@kyleconroy.com>
+- Mathieu Leplatre <mathieu@mozilla.com>
diff --git a/CHANGES b/CHANGES
index ddca03f..00c28c1 100644
--- a/CHANGES
+++ b/CHANGES
@@ -1,6 +1,11 @@
 Statsd Changelog
 ================
 
+Unreleased
+----------
+
+- Fix timer decorator with partial functions (#85)
+
 Version 3.2.2
 -------------
 
diff --git a/statsd/client.py b/statsd/client.py
index 54f62e7..6d8a921 100644
--- a/statsd/client.py
+++ b/statsd/client.py
@@ -1,6 +1,6 @@
 from __future__ import with_statement
 from collections import deque
-from functools import wraps
+import functools
 import random
 import socket
 import abc
@@ -17,6 +17,13 @@ except ImportError:
 __all__ = ['StatsClient', 'TCPStatsClient']
 
 
+def safe_wraps(wrapper, *args, **kwargs):
+    """Safely wraps partial functions."""
+    while isinstance(wrapper, functools.partial):
+        wrapper = wrapper.func
+    return functools.wraps(wrapper, *args, **kwargs)
+
+
 class Timer(object):
     """A context manager/decorator for statsd.timing()."""
 
@@ -30,7 +37,7 @@ class Timer(object):
 
     def __call__(self, f):
         """Thread-safe timing function decorator."""
-        @wraps(f)
+        @safe_wraps(f)
         def _wrapped(*args, **kwargs):
             start_time = time_now()
             try:
diff --git a/statsd/tests.py b/statsd/tests.py
index 2e388a7..b82c835 100644
--- a/statsd/tests.py
+++ b/statsd/tests.py
@@ -1,4 +1,5 @@
 from __future__ import with_statement
+import functools
 import random
 import re
 import socket
@@ -517,6 +518,18 @@ def test_timer_context_rate_tcp():
     _test_timer_context_rate(cl, 'tcp')
 
 
+def test_timer_decorator_partial_function():
+    """TCPStatsClient.timer can be used as decorator on a partial function."""
+    cl = _tcp_client()
+
+    foo = functools.partial(lambda x: x * x, 2)
+    func = cl.timer('foo')(foo)
+
+    eq_(4, func())
+
+    _timer_check(cl._sock, 1, 'tcp', 'foo', 'ms|@0.1')
+
+
 def _test_timer_decorator_rate(cl, proto):
     @cl.timer('foo', rate=0.1)
     def foo(a, b):