Give timers manual start/stop methods.

3 files changed, 63 insertions, 11 deletions

diff --git a/statsd/__init__.py b/statsd/__init__.py
index e47a620..6af74c5 100644
--- a/statsd/__init__.py
+++ b/statsd/__init__.py
@@ -16,7 +16,7 @@ except ImportError:
     settings = None
 
 from .client import StatsClient
-from ._version import __version__
+from ._version import __version__  # noqa
 
 
 __all__ = ['StatsClient', 'statsd']
diff --git a/statsd/client.py b/statsd/client.py
index 78c3a99..0770247 100644
--- a/statsd/client.py
+++ b/statsd/client.py
@@ -25,19 +25,31 @@ class Timer(object):
         return wrapper
 
     def __enter__(self):
-        self.start = time.time()
-        return self
+        return self.start()
 
     def __exit__(self, typ, value, tb):
-        dt = time.time() - self.start
-        self.ms = int(round(1000 * dt))  # Convert to ms.
+        self.stop()
+
+    def start(self):
+        self._start_time = time.time()
+        return self
+
+    def stop(self, send=True):
+        dt = time.time() - self._start_time
+        self.ms = int(round(1000 * dt))  # Convert to milliseconds.
+        if send:
+            self.send()
+        return self
+
+    def send(self):
         self.client.timing(self.stat, self.ms, self.rate)
 
 
 class StatsClient(object):
     """A client for statsd."""
 
-    def __init__(self, host='localhost', port=8125, prefix=None, maxudpsize=512):
+    def __init__(self, host='localhost', port=8125, prefix=None,
+                 maxudpsize=512):
         """Create a new client."""
         self._addr = (socket.gethostbyname(host), port)
         self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
@@ -106,7 +118,8 @@ class Pipeline(StatsClient):
         self._stats = []
 
     def _after(self, data):
-        self._stats.append(data)
+        if data is not None:
+            self._stats.append(data)
 
     def __enter__(self):
         return self
diff --git a/statsd/tests.py b/statsd/tests.py
index 75bee3d..6c7d7ad 100644
--- a/statsd/tests.py
+++ b/statsd/tests.py
@@ -18,13 +18,11 @@ def _client(prefix=None):
     return sc
 
 
-def _sock_check(cl, count, val):
+def _sock_check(cl, count, val=None):
     eq_(cl._sock.sendto.call_count, count)
-    if val:
+    if val is not None:
         val = val.encode('ascii')
         eq_(cl._sock.sendto.call_args, ((val, ADDR), {}))
-    else:
-        eq_(cl._sock.sendto.call_args, None)
 
 
 class assert_raises(object):
@@ -288,6 +286,37 @@ def test_timer_decorator_exceptions():
     _timer_check(sc, 1, 'foo', 'ms')
 
 
+def test_imperative_timer():
+    sc = _client()
+
+    t = sc.timer('foo').start()
+    t.stop()
+
+    _timer_check(sc, 1, 'foo', 'ms')
+
+
+def test_imperative_timer_no_send():
+    sc = _client()
+
+    t = sc.timer('foo').start()
+    t.stop(send=False)
+    _sock_check(sc, 0)
+
+    t.send()
+    _timer_check(sc, 1, 'foo', 'ms')
+
+
+@mock.patch.object(random, 'random', lambda: -1)
+def test_imperative_timer_rate():
+    sc = _client()
+
+    t = sc.timer('foo', rate=0.5)
+    t.start()
+    t.stop()
+
+    _timer_check(sc, 1, 'foo', 'ms@0.5')
+
+
 def test_pipeline():
     sc = _client()
     pipe = sc.pipeline()
@@ -303,6 +332,7 @@ def test_pipeline_null():
     sc = _client()
     pipe = sc.pipeline()
     pipe.send()
+    _sock_check(sc, 0)
 
 
 def test_pipeline_manager():
@@ -332,6 +362,15 @@ def test_pipeline_timer_decorator():
     _timer_check(sc, 1, 'foo', 'ms')
 
 
+def test_pipeline_timer_imperative():
+    sc = _client()
+    with sc.pipeline() as pipe:
+        t = pipe.timer('foo').start()
+        t.stop()
+        _sock_check(sc, 0)
+    _timer_check(sc, 1, 'foo', 'ms')
+
+
 def test_pipeline_empty():
     """Pipelines should be empty after a send() call."""
     sc = _client()