Merge pull request #7 from trehn/feature/contextmanager

Expose a context manager for histograms (continued)

README.rst

@@ -54,7 +54,7 @@ an internal registry, so you can access it in different places in your applicati
 
 The ``metrics`` registry is thread-safe, you can safely use it in multi-threaded web servers.
 
-Now let's decorate some function::
+Using the ``with_histogram`` decorator we can time a function::
 
     >>> import time, random
     >>> @metrics.with_histogram("test")
@@ -70,6 +70,14 @@ and let's see the results::
     >>> metrics.get("test")
     {'arithmetic_mean': 0.41326093673706055, 'kind': 'histogram', 'skewness': 0.2739718270714368, 'harmonic_mean': 0.14326954591313346, 'min': 0.0613858699798584, 'standard_deviation': 0.4319169569113129, 'median': 0.2831099033355713, 'histogram': [(1.0613858699798584, 3), (2.0613858699798584, 0)], 'percentile': [(50, 0.2831099033355713), (75, 0.2831099033355713), (90, 0.895287036895752), (95, 0.895287036895752), (99, 0.895287036895752), (99.9, 0.895287036895752)], 'n': 3, 'max': 0.895287036895752, 'variance': 0.18655225766752892, 'geometric_mean': 0.24964828731906127, 'kurtosis': -2.3333333333333335}
 
+It is also possible to time specific sections of the code by using the ``timer`` context manager::
+
+    >>> import time, random
+    ... def my_worker():
+    ...     with metrics.timer("test"):
+    ...         time.sleep(random.random())
+    ...
+
 Let's print the metrics data on the screen every 5 seconds::
 
     >>> from appmetrics import reporter

appmetrics/histogram.py

@@ -32,15 +32,16 @@
 
 def search_greater(values, target):
     """
-    Return the first index for which target is greater or equal to the first item of the tuple found in values
+    Return the first index for which target is greater or equal to the first
+    item of the tuple found in values
     """
     first = 0
     last = len(values)
 
     while first < last:
-        middle = (first+last)//2
+        middle = (first + last) // 2
         if values[middle][0] < target:
-            first = middle+1
+            first = middle + 1
         else:
             last = middle
 
@@ -51,13 +52,15 @@ class ReservoirBase(object):
     __metaclass__ = abc.ABCMeta
 
     """
-    Base class for reservoirs. Subclass and override _do_add, _get_values and _same_parameters
+    Base class for reservoirs. Subclass and override _do_add, _get_values and
+    _same_parameters
     """
 
     def add(self, value):
         """
         Add a value to the reservoir
-        The value will be casted to a floating-point, so a TypeError or a ValueError may be raised.
+        The value will be casted to a floating-point, so a TypeError or a
+        ValueError may be raised.
         """
 
         if not isinstance(value, float):
@@ -73,7 +76,6 @@ def values(self):
 
         return self._get_values()
 
-
     @property
     def sorted_values(self):
         """
@@ -84,7 +86,8 @@ def sorted_values(self):
 
     def same_kind(self, other):
         """
-        Return True if "other" is an object of the same type and it was instantiated with the same parameters
+        Return True if "other" is an object of the same type and it was
+        instantiated with the same parameters
         """
 
         return type(self) is type(other) and self._same_parameters(other)
@@ -104,15 +107,17 @@ def _get_values(self):
     @abc.abstractmethod
     def _same_parameters(self, other):
         """
-        Return True if this object has been instantiated with the same parameters as "other".
+        Return True if this object has been instantiated with the same
+        parameters as "other".
         Override in subclasses
         """
 
 
 class UniformReservoir(ReservoirBase):
     """
-    A random sampling reservoir of floating-point values. Uses Vitter's Algorithm R to produce a statistically
-    representative sample (http://www.cs.umd.edu/~samir/498/vitter.pdf)
+    A random sampling reservoir of floating-point values. Uses Vitter's
+    Algorithm R to produce a statistically representative sample
+    (http://www.cs.umd.edu/~samir/498/vitter.pdf)
     """
 
     def __init__(self, size=DEFAULT_UNIFORM_RESERVOIR_SIZE):
@@ -129,8 +134,8 @@ def _do_add(self, value):
                 self._values[self.count] = value
                 changed = True
             else:
-                # not randint() because it yields different values on python 3, it
-                # would be a nightmare to test.
+                # not randint() because it yields different values on
+                # python 3, it would be a nightmare to test.
                 k = int(random.uniform(0, self.count))
                 if k < self.size:
                     self._values[k] = value
@@ -200,7 +205,8 @@ def _do_add(self, value):
     def tick(self, now):
         target = now - self.window_size
 
-        # the values are sorted by the first element (timestamp), so let's perform a dichotomic search
+        # the values are sorted by the first element (timestamp), so let's
+        # perform a dichotomic search
         idx = search_greater(self._values, target)
 
         # older values found, discard them
@@ -229,14 +235,16 @@ class ExponentialDecayingReservoir(ReservoirBase):
     See http://dimacs.rutgers.edu/~graham/pubs/papers/fwddecay.pdf
     """
 
-    #TODO: replace the sort()s with a proper data structure (btree/skiplist). However, since the list
-    # is keep sorted (and it should be very small), the sort() shouldn't dramatically slow down
-    # the insertions, also considering that the search can be log(n) in that way
+    # TODO: replace the sort()s with a proper data structure (btree/skiplist).
+    # However, since the list is keep sorted (and it should be very small),
+    # the sort() shouldn't dramatically slow down the insertions, also
+    # considering that the search can be log(n) in that way
 
     RESCALE_THRESHOLD = 3600
     EPSILON = 1e-12
 
-    def __init__(self, size=DEFAULT_UNIFORM_RESERVOIR_SIZE, alpha=DEFAULT_EXPONENTIAL_DECAY_FACTOR):
+    def __init__(self, size=DEFAULT_UNIFORM_RESERVOIR_SIZE,
+                 alpha=DEFAULT_EXPONENTIAL_DECAY_FACTOR):
         self.size = size
         self.alpha = alpha
         self.start_time = time.time()
@@ -249,13 +257,15 @@ def __init__(self, size=DEFAULT_UNIFORM_RESERVOIR_SIZE, alpha=DEFAULT_EXPONENTIA
 
     def _lookup(self, timestamp):
         """
-        Return the index of the value associated with "timestamp" if any, else None.
-        Since the timestamps are floating-point values, they are considered equal if their absolute
-        difference is smaller than self.EPSILON
+        Return the index of the value associated with "timestamp" if any, else
+        None. Since the timestamps are floating-point values, they are
+        considered equal if their absolute difference is smaller than
+        self.EPSILON
         """
 
         idx = search_greater(self._values, timestamp)
-        if idx < len(self._values) and math.fabs(self._values[idx][0] - timestamp) < self.EPSILON:
+        if (idx < len(self._values)
+                and math.fabs(self._values[idx][0] - timestamp) < self.EPSILON):
             return idx
         return None
 
@@ -274,7 +284,7 @@ def _do_add(self, value):
         self.rescale(now)
 
         rnd = random.random()
-        weighted_time = self.weight(now-self.start_time) / rnd
+        weighted_time = self.weight(now - self.start_time) / rnd
 
         changed = False
 
@@ -305,14 +315,13 @@ def rescale(self, now):
                 original_values = self._values[:]
                 self._values = []
                 for i, (k, v) in enumerate(original_values):
-                    k *= math.exp(-self.alpha * (now-self.start_time))
+                    k *= math.exp(-self.alpha * (now - self.start_time))
                     self._put(k, v)
 
                 self.count = len(self._values)
                 self.start_time = now
                 self.next_scale_time = self.start_time + self.RESCALE_THRESHOLD
 
-
     def _get_values(self):
         return [y for x, y in self._values[:max(self.count, self.size)]]
 
@@ -324,7 +333,8 @@ def __repr__(self):
 
 
 class Histogram(object):
-    """A metric which calculates some statistics over the distribution of some values"""
+    """A metric which calculates some statistics over the distribution of some
+    values"""
 
     def __init__(self, reservoir):
         self.reservoir = reservoir

appmetrics/metrics.py

@@ -18,6 +18,7 @@
 Main interface module
 """
 
+from contextlib import contextmanager
 import functools
 import threading
 import time
@@ -160,13 +161,12 @@ def new_reservoir(reservoir_type='uniform', *reservoir_args, **reservoir_kwargs)
     return reservoir_cls(*reservoir_args, **reservoir_kwargs)
 
 
-def with_histogram(name, reservoir_type="uniform", *reservoir_args, **reservoir_kwargs):
+def get_or_create_histogram(name, reservoir_type, *reservoir_args, **reservoir_kwargs):
     """
-    Time-measuring decorator: the time spent in the wrapped function is measured
-    and added to the named metric.
-    metric_args and metric_kwargs are passed to new_histogram()
+    Will return a histogram matching the given parameters or raise
+    DuplicateMetricError if it can't be created due to a name collision
+    with another histogram with different parameters.
     """
-
     reservoir = new_reservoir(reservoir_type, *reservoir_args, **reservoir_kwargs)
 
     try:
@@ -181,6 +181,18 @@ def with_histogram(name, reservoir_type="uniform", *reservoir_args, **reservoir_
             raise DuplicateMetricError(
                 "Metric {!r} already exists with a different reservoir: {}".format(name, hmetric.reservoir))
 
+    return hmetric
+
+
+def with_histogram(name, reservoir_type="uniform", *reservoir_args, **reservoir_kwargs):
+    """
+    Time-measuring decorator: the time spent in the wrapped function is measured
+    and added to the named metric.
+    metric_args and metric_kwargs are passed to new_histogram()
+    """
+
+    hmetric = get_or_create_histogram(name, reservoir_type, *reservoir_args, **reservoir_kwargs)
+
     def wrapper(f):
 
         @functools.wraps(f)
@@ -229,6 +241,21 @@ def fun(*args, **kwargs):
     return wrapper
 
 
+@contextmanager
+def timer(name, reservoir_type="uniform", *reservoir_args, **reservoir_kwargs):
+    """
+    Time-measuring context manager: the time spent in the wrapped block
+    if measured and added to the named metric.
+    """
+
+    hmetric = get_or_create_histogram(name, reservoir_type, *reservoir_args, **reservoir_kwargs)
+
+    t1 = time.time()
+    yield
+    t2 = time.time()
+    hmetric.notify(t2 - t1)
+
+
 def tag(name, tag_name):
     """
     Tag the named metric with the given tag.

appmetrics/tests/test_metrics.py

@@ -372,6 +372,38 @@ def f2(v1, v2):
 
             return v1*v2
 
+    @mock.patch('appmetrics.histogram.Histogram.notify')
+    def test_timer(self, notify):
+        with mm.timer("test"):
+            pass
+
+        assert_equal(notify.call_count, 1)
+
+    @mock.patch('appmetrics.histogram.Histogram.notify')
+    def test_timer_multiple(self, notify):
+        with mm.timer("test"):
+            pass
+
+        with mm.timer("test"):
+            pass
+
+        assert_equal(notify.call_count, 2)
+
+    @raises(exceptions.DuplicateMetricError)
+    def test_timer_multiple_different_reservoir(self):
+        with mm.timer("test", reservoir_type="sliding_window"):
+            pass
+
+        with mm.timer("test"):
+            pass
+
+    @raises(exceptions.DuplicateMetricError)
+    def test_timer_multiple_different_type(self):
+        mm.new_gauge("test")
+
+        with mm.timer("test"):
+            pass
+
     @raises(exceptions.InvalidMetricError)
     def test_tag_invalid_name(self):
         mm.tag("test", "test")