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

import asyncio

from contextlib import contextmanager
from functools import wraps
from sentry_sdk import Hub


# https://stackoverflow.com/questions/44169998/how-to-create-a-python-decorator-that-can-wrap-either-coroutine-or-function
def span(op, desc_or_func=None):
    def inner(func):
        @contextmanager
        def wrap_with_span(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):
                yield

        @wraps(func)
        def wrapper(*args, **kwargs):
            if not asyncio.iscoroutinefunction(func):
                with wrap_with_span(args, kwargs):
                    return func(*args, **kwargs)
            else:

                async def tmp():
                    with wrap_with_span(args, kwargs):
                        return await func(*args, **kwargs)

                return tmp()

        return wrapper

    return inner

Fix span decorator to work with asyncio
Fix span decorator to work with asyncio

+ import asyncio
+ 
+ from contextlib import contextmanager
  from functools import wraps
  from sentry_sdk import Hub
  
  
+ # https://stackoverflow.com/questions/44169998/how-to-create-a-python-decorator-that-can-wrap-either-coroutine-or-function
  def span(op, desc_or_func=None):
      def inner(func):
-         @wraps(func)
+         @contextmanager
-         def wrapped(*args, **kwargs):
+         def wrap_with_span(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)
+                 yield
  
+         @wraps(func)
+         def wrapper(*args, **kwargs):
+             if not asyncio.iscoroutinefunction(func):
+                 with wrap_with_span(args, kwargs):
+                     return func(*args, **kwargs)
+             else:
+ 
+                 async def tmp():
+                     with wrap_with_span(args, kwargs):
+                         return await func(*args, **kwargs)
+ 
+                 return tmp()
+ 
-         return wrapped
+         return wrapper
  
      return inner
  
Fix span decorator to work with asyncio
## Code Before:
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

## Instruction:
Fix span decorator to work with asyncio
## Code After:
import asyncio

from contextlib import contextmanager
from functools import wraps
from sentry_sdk import Hub


# https://stackoverflow.com/questions/44169998/how-to-create-a-python-decorator-that-can-wrap-either-coroutine-or-function
def span(op, desc_or_func=None):
    def inner(func):
        @contextmanager
        def wrap_with_span(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):
                yield

        @wraps(func)
        def wrapper(*args, **kwargs):
            if not asyncio.iscoroutinefunction(func):
                with wrap_with_span(args, kwargs):
                    return func(*args, **kwargs)
            else:

                async def tmp():
                    with wrap_with_span(args, kwargs):
                        return await func(*args, **kwargs)

                return tmp()

        return wrapper

    return inner
