Merge pull request #36 from ooda/profile-decorator

Add two decorators: @Profile and @line_profile.

cloudly/decorators.py

@@ -6,6 +6,8 @@
 import inspect
 from functools import wraps, partial
 from datetime import datetime, timedelta
+import cProfile
+import pstats
 
 
 class Memoized(object):
@@ -123,3 +125,44 @@ def burst_generator(length, generator):
         if len(burst) >= length:
             yield burst
             burst = []
+
+
+def profile(sort="tottime", return_pstats=False):
+    def decorator(fn):
+        @wraps(fn)
+        def wrapped_fn(*args, **kwargs):
+            profile = cProfile.Profile()
+            try:
+                profile.enable()
+                result = fn(*args, **kwargs)
+                profile.disable()
+                ps = pstats.Stats(profile)
+                if return_pstats:
+                    return ps
+                else:
+                    return result
+            finally:
+                if not return_pstats:
+                    ps.sort_stats(sort)
+                    ps.print_stats()
+        return wrapped_fn
+    return decorator
+
+
+def line_profile(follow=[]):
+    from line_profiler import LineProfiler
+
+    def decorator(fn):
+        @wraps(fn)
+        def wrapped_fn(*args, **kwargs):
+            try:
+                profiler = LineProfiler()
+                profiler.add_function(fn)
+                for f in follow:
+                    profiler.add_function(f)
+                profiler.enable_by_count()
+                return fn(*args, **kwargs)
+            finally:
+                profiler.print_stats()
+        return wrapped_fn
+    return decorator

tests/unittests/cloudly/test_decorators.py

@@ -1,7 +1,7 @@
 from time import sleep
 from datetime import datetime, timedelta
 
-from cloudly.decorators import burst, throttle
+from cloudly.decorators import burst, throttle, profile, line_profile
 
 
 def test_burst():
@@ -67,3 +67,31 @@ def method(self, arg):
         else:
             assert result is None
         sleep(interval)
+
+
+def test_profile():
+    def get_number():
+        for x in xrange(5000000):
+            yield x
+
+    @profile()
+    def expensive_function():
+        for x in get_number():
+            i = x ^ x ^ x
+        return i
+
+    expensive_function()
+
+
+def test_line_profile():
+    def get_number():
+        for x in xrange(50000):
+            yield x
+
+    @line_profile(follow=[get_number])
+    def expensive_function():
+        for x in get_number():
+            i = x ^ x ^ x
+        return i
+
+    expensive_function()