Allow trace marker to be use as decorator (#325)

Closes #297

README.md

@@ -390,17 +390,32 @@ Or you can use it as a context manager:
 
 ```python
 from iopipe import IOpipe
-from iopipe.contrib.trace import TracePlugin
 
-iopipe = IOpipe(plugins=[TracePlugin()])
+iopipe = IOpipe()
 
 @iopipe
 def handler(event, context):
     with context.iopipe.mark('expensive operation'):
         # do something here
 ```
 
-Any block of code wrapped with `start` and `end` or using the context manager will be traced and the data collected will be available on your IOpipe dashboard.
+Or you can use it as a decorator:
+
+```python
+from iopipe import IOpipe
+
+iopipe = IOpipe()
+
+@iopipe
+def handler(event, context):
+    @context.iopipe.mark.decorator('expensive operation'):
+	def expensive_operation():
+        # do something here
+
+    expensive_operation()
+```
+
+Any block of code wrapped with `start` and `end` or using the context manager or decorator will be traced and the data collected will be available on your IOpipe dashboard.
 
 By default, the trace plugin will auto-measure any trace you make. But you can disable this by setting `auto_measure` to `False`:
 

iopipe/contrib/trace/marker.py

@@ -1,3 +1,6 @@
+import functools
+
+
 class Marker(object):
     def __init__(self, timeline, context):
         self.timeline = timeline
@@ -18,6 +21,9 @@ def __exit__(self, type, value, traceback):
             self.end(self.contexts[-1])
             self.contexts.pop()
 
+    def decorator(self, name):
+        return MarkerDecorator(self, name)
+
     def start(self, name):
         self.timeline.mark("start:%s" % name)
         self.context.iopipe.label("@iopipe/plugin-trace")
@@ -49,3 +55,22 @@ def http_trace(self, trace, request, response):
             entry["response"] = response
 
         self.context.instance.report.http_trace_entries.append(entry)
+
+
+class MarkerDecorator(object):
+    def __init__(self, marker, name):
+        self.marker = marker
+        self.name = name
+
+    def __call__(self, func):
+        self.marker.contexts.append(self.name)
+
+        @functools.wraps(func)
+        def wrapped(*args, **kwargs):
+            self.marker.start(self.marker.contexts[-1])
+            return_value = func(*args, **kwargs)
+            self.marker.end(self.marker.contexts[-1])
+            self.marker.contexts.pop()
+            return return_value
+
+        return wrapped

tests/contrib/trace/test_marker.py

@@ -46,3 +46,13 @@ def test_marker__nested_context_manager(marker):
             pass
 
     assert len(marker.timeline.get_entries()) == 4
+
+
+def test_marker__decorator(marker):
+    @marker.decorator("foobar")
+    def mock_func():
+        pass
+
+    mock_func()
+
+    assert len(marker.timeline.get_entries()) == 2