Function annotate_stack

axlearn/common/traceback_util.py:267–298

A decorator that creates a wrapper around `fn` that annotates the call with annotations given by the supplied keyword arguments. These annotations can be retrieved from a traceback using `walk_annotated_tb()`. Example: ``` fn = lambda x: x annotated_fn = annotat

(**aux)

# function, there is no way of getting the function object itself that correctly handles anonymous
# functions. Otherwise, we could just set an attribute on the function.
def annotate_stack(**aux) -> Callable:
    """A decorator that creates a wrapper around `fn` that annotates the call with annotations given
    by the supplied keyword arguments.

    These annotations can be retrieved from a traceback using `walk_annotated_tb()`.

    Example:
        ```
        fn = lambda x: x
        annotated_fn = annotate_stack(my_annotation="Hello, world!")(fn)
        ```

    Args:
        aux: The auxiliary data with which to annotate calls to the wrapped function.

    Returns:
        A wrapped function that calls `fn` after annotating the call with `aux`.
    """

    if not is_stack_summary_enabled():
        return lambda fn: fn

    def decorator(fn: Callable) -> Callable:
        @functools.wraps(fn)
        def stack_annotation_wrapper(*args, **kwargs):
            # Retrieved via `frame.f_locals["aux_"]` by `_walk_stack` below.
            aux_ = aux  # noqa: F841
            return fn(*args, **kwargs)

        return stack_annotation_wrapper

    return decorator


def _is_annotation_frame(frame: types.FrameType) -> bool: