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)]
