write and apply step_log_wrap for functions that return values

1st attempt hit a roadblock attempted to apply following wrapper to test_problem.model_state.ModelState.apply_precond_jacobian as

    @res_step_log_wrap(
        step="apply_precond_jacobian complete for {res_fname}", res_type=ModelState
    )
    # pylint: disable=unused-argument
    def apply_precond_jacobian(self, precond_fname, res_fname, solver_state):

This generates a NameError when run:

  File "/Users/klindsay/codes/Newton-Krylov_OOC/src/test_problem/model_state.py", line 329, in ModelState
    step="apply_precond_jacobian complete for {res_fname}", res_type=ModelState
NameError: name 'ModelState' is not defined

Replacing ModelState with type(self) also generates a NameError:

NameError: name 'self' is not defined

I don't see how to get ModelState into res_step_log_wrap.

wrapper is

def res_step_log_wrap(step, res_type, per_iteration=True, post_exit=False):
    """
    Decorator for wrapping functions with args inside step_logged/log_step checks.
    It is for functions that return a model state from res_fname.
    res_fname and solver_state are assumed to be keyword arguments to func.
    step is the string argument getting passed to sover_state methods.
    Formatting using .format is applied to step, using the keyword arguments of func,
    to enable step to depend on func's arguments.
    """

    def outer_wrapper(func):
        @functools.wraps(func)  # to propagate metadata from func through wrapper
        def inner_wrapper(*args, **kwargs):
            solver_state = kwargs["solver_state"]
            if solver_state is not None:
                if solver_state.step_logged(step.format(**kwargs), per_iteration):
                    return res_type(kwargs["res_fname"])
            res = func(*args, **kwargs)
            if solver_state is not None:
                solver_state.log_step(step.format(**kwargs), per_iteration)
            if post_exit:
                raise SystemExit
            return res

        return inner_wrapper

    return outer_wrapper

klindsay28 / Newton-Krylov_OOC

write and apply step_log_wrap for functions that return values #34