def named_call(class_fn):
"""Labels a method for labelled traces in profiles."""
# NB: due to the ordering of method decorators, we must re-wrap the class_fn
# to maintain Module state correctly for multiple invocations. If we want to
# save another stacktrace entry we could instead replicate its logic below.
rewrapped_fn = wrap_method(class_fn)
@functools.wraps(class_fn)
def wrapped_fn(self, *args, **kwargs):
fn_name = class_fn.__name__
method_suffix = f'.{fn_name}' if fn_name != '__call__' else ''
module_name = self.name or self.__class__.__name__
full_name = f'{module_name}{method_suffix}'
# make a scope-function to transform
def core_fn(scopes, *args, **kwargs):
cloned = set_module_scopes(self, scopes)
cloned._state = copy.deepcopy(self._state) # pylint: disable=protected-access
res = rewrapped_fn(cloned, *args, **kwargs)
# preserve submodule-tree stripped of scopes/tracers for introspection
object.__setattr__(self, 'children', clean_clone(cloned).children)
self._state = copy.deepcopy(cloned._state) # pylint: disable=protected-access
return res
# here we apply the given lifting transform to the scope-ingesting fn
trafo_fn = lift.named_call(core_fn, full_name)
return trafo_fn(get_module_scopes(self), *args, **kwargs)
return wrapped_fn
def decorator_lift_transform(transform, class_fn, *trafo_args, **trafo_kwargs):
# NB: due to the ordering of method decorators, we must re-wrap the class_fn
# to maintain Module state correctly for multiple invocations. If we want to
# save another stacktrace entry we could instead replicate its logic below.
rewrapped_fn = wrap_method(class_fn)
@functools.wraps(class_fn)
def wrapped_fn(self, *args, **kwargs):
# make a scope-function to transform
def core_fn(scopes, *args, **kwargs):
cloned = set_module_scopes(self, scopes)
cloned._state = copy.deepcopy(self._state) # pylint: disable=protected-access
res = rewrapped_fn(cloned, *args, **kwargs)
self._state = copy.deepcopy(cloned._state) # pylint: disable=protected-access
return res
# here we apply the given lifting transform to the scope-ingesting fn
trafo_fn = transform(core_fn, *trafo_args, **trafo_kwargs)
return trafo_fn(get_module_scopes(self), *args, **kwargs)
return wrapped_fn
def decorator_lift_transform(transform, class_fn, *trafo_args, **trafo_kwargs):
# NB: due to the ordering of method decorators, we must re-wrap the class_fn
# to maintain Module state correctly for multiple invocations. If we want to
# save another stacktrace entry we could instead replicate its logic below.
rewrapped_fn = wrap_method(class_fn)
@functools.wraps(class_fn)
def wrapped_fn(self, *args, **kwargs):
# make a scope-function to transform
def core_fn(scopes, *args, **kwargs):
cloned = set_module_scopes(self, scopes)
res = rewrapped_fn(cloned, *args, **kwargs)
# preserve submodule-tree stripped of scopes/tracers for introspection
object.__setattr__(self, 'children', clean_clone(cloned).children)
return res
# here we apply the given lifting transform to the scope-ingesting fn
trafo_fn = transform(core_fn, *trafo_args, **trafo_kwargs)
return trafo_fn(get_module_scopes(self), *args, **kwargs)
return wrapped_fn
