diff options
| -rw-r--r-- | AUTHORS | 1 | ||||
| -rw-r--r-- | CHANGES | 5 | ||||
| -rw-r--r-- | statsd/client.py | 11 | ||||
| -rw-r--r-- | statsd/tests.py | 13 |
4 files changed, 28 insertions, 2 deletions
@@ -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> @@ -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): |