def set_trace(func: typing.Callable, trace_func: typing.Callable): """Set `trace_func` to be called at the beginning of `func`.""" func = _unwrap(func) trace_fn_name = _get_trace_fn_name(func) trace_func = partial(trace_func, event="call", arg=None) _update_globals(func.__globals__, trace_fn_name, trace_func, cmp_=lambda p: p.func) if not hasattr(func, _original_code_attribute_name): setattr(func, _original_code_attribute_name, func.__code__) func.__code__ = _make_code(func, trace_fn_name) _update_globals(func.__globals__, trace_fn_name, trace_func) _update_globals(func.__globals__, _get_frame_fn_name, sys._getframe)
def unset_trace(func: typing.Callable): """Remove trace function from `func`.""" if hasattr(func, _original_code_attribute_name): func.__code__ = getattr(func, _original_code_attribute_name) delattr(func, _original_code_attribute_name) del func.__globals__[_get_trace_fn_name(func)]