def named_call(class_fn):
"""Labels a method for labelled traces in profiles."""
# Due to the ordering of method decorators, we must wrap the class_fn
# with the module state management wrapper first to maintain Module state correctly.
prewrapped_fn = wrap_method_once(class_fn)
@functools.wraps(prewrapped_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)
object.__setattr__(cloned, '_state', self._state.export()) # pylint: disable=protected-access
res = prewrapped_fn(cloned, *args, **kwargs)
self._state.reimport(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):
# Due to the ordering of method decorators, we must wrap the class_fn
# with the module state management wrapper first to maintain Module state correctly.
prewrapped_fn = wrap_method_once(class_fn)
@functools.wraps(prewrapped_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 = self._state.export() # pylint: disable=protected-access
res = prewrapped_fn(cloned, *args, **kwargs)
self._state.reimport(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
