def wrap_lazy(
    obj: Any,
    on_start: Optional[Callable[[], None]] = None,
    wrap: Optional[Callable[[T], T]] = None,
    on_done: Optional[Callable[[Any], Any]] = None,  # Any may be T or List[T]
    context_vars: Optional[ContextVarsOrValues] = None,
) -> Any:
    """Wrap a lazy value in one that will call
    callbacks at various points in the generation process.

    Args:
        obj: The lazy value.

        on_start: The callback to call when the wrapper is created.

        wrap: The callback to call with the result of each iteration of the
            wrapped generator or the result of an awaitable. This should return
            the value or a wrapped version.

        on_done: The callback to call when the wrapped generator is exhausted or
            awaitable is ready.

        context_vars: The context variables to copy over to the wrapped
            generator. If None, all context variables are taken with their
            present values. See
            [with_context][trulens.core.utils.python.with_context].
    """

    if not WRAP_LAZY:
        return obj

    if inspect.isasyncgen(obj):
        return wrap_async_generator(
            obj,
            on_iter=on_start,
            wrap=wrap,
            on_done=on_done,
            context_vars=context_vars,
        )

    if inspect.isgenerator(obj):
        return wrap_generator(
            obj,
            on_iter=on_start,
            wrap=wrap,
            on_done=on_done,
            context_vars=context_vars,
        )

    if inspect.isawaitable(obj):
        return wrap_awaitable(
            obj,
            on_await=on_start,
            wrap=wrap,
            on_done=on_done,
            context_vars=context_vars,
        )

    raise ValueError(f"Object of type {type(obj)} is not lazy.")
Wrap a lazy value in one that will call
    callbacks at various points in the generation process.

    Args:
        obj: The lazy value.

        on_start: The callback to call when the wrapper is created.

        wrap: The callback to call with the result of each iteration of the
            wrapped generator or the result of an awaitable. This should return
            the value or a wrapped version.

        on_done: The callback to call when the wrapped generator is exhausted or
            awaitable is ready.

        context_vars: The context variables to copy over to the wrapped
            generator. If None, all context variables are taken with their
            present values. See
            [with_context][trulens.core.utils.python.with_context].
    
src/core/trulens/core/utils/python.py