from functools import wraps
from sentry_sdk import Hub


def span(op, desc_or_func=None):
    def inner(func):
        @wraps(func)
        def wrapped(*args, **kwargs):
            if callable(desc_or_func):
                description = desc_or_func(*args, **kwargs)
            else:
                description = desc_or_func
            with Hub.current.start_span(op=op, description=description):
                return func(*args, **kwargs)

        return wrapped

    return inner

@@ -1,12 +1,66 @@
+import asyncio%0A%0Afrom contextlib import contextmanager%0A
 from functoo
@@ -104,16 +104,139 @@
 t Hub%0A%0A%0A
+# https://stackoverflow.com/questions/44169998/how-to-create-a-python-decorator-that-can-wrap-either-coroutine-or-function%0A
 def span
@@ -294,19 +294,22 @@
    @
-wraps(func)
+contextmanager
 %0A   
@@ -325,21 +325,25 @@
 wrap
-ped(*
+_with_span(
 args, 
-**
 kwar
@@ -600,35 +600,422 @@
     
-return func(*args, **kwargs
+yield%0A%0A        @wraps(func)%0A        def wrapper(*args, **kwargs):%0A            if not asyncio.iscoroutinefunction(func):%0A                with wrap_with_span(args, kwargs):%0A                    return func(*args, **kwargs)%0A            else:%0A%0A                async def tmp():%0A                    with wrap_with_span(args, kwargs):%0A                        return await func(*args, **kwargs)%0A%0A                return tmp(
 )%0A%0A 
@@ -1034,17 +1034,17 @@
 n wrappe
-d
+r
 %0A%0A    re

Fix span decorator to work with asyncio