From cadd2d5f7cdc904431daa67adea08cfe29552378 Mon Sep 17 00:00:00 2001
From: Michael Lavers <kolanos@gmail.com>
Date: Tue, 11 Jun 2019 08:55:32 -0700
Subject: [PATCH] Allow trace marker to be use as decorator (#325)

Closes #297
---
 README.md                          | 21 ++++++++++++++++++---
 iopipe/contrib/trace/marker.py     | 25 +++++++++++++++++++++++++
 tests/contrib/trace/test_marker.py | 10 ++++++++++
 3 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 9743fdef..44830f9e 100644
--- a/README.md
+++ b/README.md
@@ -390,9 +390,8 @@ 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):
@@ -400,7 +399,23 @@ def handler(event, context):
         # 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`:
 
diff --git a/iopipe/contrib/trace/marker.py b/iopipe/contrib/trace/marker.py
index 26f64f75..91520c71 100644
--- a/iopipe/contrib/trace/marker.py
+++ b/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
diff --git a/tests/contrib/trace/test_marker.py b/tests/contrib/trace/test_marker.py
index 74ce64e5..31679a0c 100644
--- a/tests/contrib/trace/test_marker.py
+++ b/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