def wrap_awaitable(
    awaitable: Awaitable[T],
    on_await: Optional[Callable[[], Any]] = None,
    wrap: Optional[Callable[[T], T]] = None,
    on_done: Optional[Callable[[T], T]] = None,
    context_vars: Optional[ContextVarsOrValues] = None,
) -> Awaitable[T]:
    """Wrap an awaitable in another awaitable that will call callbacks before
    and after the given awaitable finishes.

    !!! Important
        This method captures a [Context][contextvars.Context] at the time this
        method is called and copies it over to the wrapped awaitable.

    Note that the resulting awaitable needs to be awaited for the callback to
    eventually trigger.

    Args:
        awaitable: The awaitable to wrap.

        on_await: The callback to call when the wrapper awaitable is awaited but
            before the wrapped awaitable is awaited.

        wrap: The callback to call with the result of the wrapped awaitable
            once it is ready. This should return the value or a wrapped version.

        on_done: For compatibility with generators, this is called after wrap.

        context_vars: The context variables to copy over to the wrapped
            awaitable. If None, all context variables are copied. See
            [with_context][trulens.core.utils.python.with_context].
    """

    async def wrapper(awaitable):
        async with awith_context(context_vars):
            if on_await is not None:
                on_await()

            val = await awaitable

            if wrap is not None:
                val = wrap(val)  # allow handlers to transform the value

            if on_done is not None:
                val = on_done(val)

            return val

    wrapper.__name__ = f"tru_wrapped_{awaitable.__class__.__name__}"

    return wrapper(awaitable)
Wrap an awaitable in another awaitable that will call callbacks before
    and after the given awaitable finishes.

    !!! Important
        This method captures a [Context][contextvars.Context] at the time this
        method is called and copies it over to the wrapped awaitable.

    Note that the resulting awaitable needs to be awaited for the callback to
    eventually trigger.

    Args:
        awaitable: The awaitable to wrap.

        on_await: The callback to call when the wrapper awaitable is awaited but
            before the wrapped awaitable is awaited.

        wrap: The callback to call with the result of the wrapped awaitable
            once it is ready. This should return the value or a wrapped version.

        on_done: For compatibility with generators, this is called after wrap.

        context_vars: The context variables to copy over to the wrapped
            awaitable. If None, all context variables are copied. See
            [with_context][trulens.core.utils.python.with_context].
    
src/core/trulens/core/utils/python.py