Exception cleanup

[liamhuber]

contextlib.ExitStack provides a with context onto which we can slap callbacks for when the context is exited. However, these all always fire when the context is exited.

Sometimes in pyiron_workflow, I have a series of operations that make progressive modifications to some state, but if I encounter an error I want to go back and revert whatever changes have been made so far before raising the error. Right now I manually handle these callbacks, but It seems to me like this is a special case of ExitStack, "ExceptionStack" maybe.

This can be pretty easily constructed as a custom context manager and I thought it might be a nice addition here.

Feeback from anyone/@pmrv?

[pmrv]

Can't you wrap your callbacks as contextmanagers and push those onto the exit stack?

@contextlib.contextmanager
def onerror(func):
  try:
    yield
  except:
    func()
    raise


# somewhere in the application code
with Exitstack() as stack:
  stack.enter_context(onerror(rollback1))
  ...
  stack.enter_context(onerror(rollback2))
  ...

[liamhuber]

Yes, 100%. There's three things I like less about this approach:

1. In my use cases my stack is only for when an exception is encountered, so this winds up being more typing because I need the stack plus the wrapper for each callback
2. It requires a (subjectively) less intuitive method call (.enter_context vs .callback)
3. It has a nasty trap if you accidentally use it with the regular .callback

import contextlib

@contextlib.contextmanager
def onerror(func, *args, **kwargs):
  try:
    yield
  except:
    func(*args, **kwargs)
    raise


def rollback(msg):
    print("\tROLLBACK:", msg)

def callback(msg):
    print("\tCALLBACK:", msg)

print("Without Error")
with contextlib.ExitStack() as stack:
    stack.callback(callback, "via .callback")
    stack.enter_context(onerror(rollback, "via .enter_context"))

print("With Error")
with contextlib.ExitStack() as stack:
    stack.callback(callback, "via .callback")
    stack.enter_context(onerror(rollback, "via .enter_context"))

    # Potentially annoying gotcha:
    stack.callback(onerror(rollback), "via .callback")  # Never seen
    # Mixing the wrapper with the usual callback will fail,
    # because .callback is expecting a function,
    # but onerror is returning a generator
    raise RuntimeError("Application error")

Without Error
	CALLBACK: via .callback
With Error
	ROLLBACK: via .enter_context
	CALLBACK: via .callback

Although my use case is for callbacks that only happen on exception, let's imagine the mixed case. For my proposal of a separate stack, that means having two with clauses.

import contextlib

class ExceptionExitStack(contextlib.ExitStack):
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:  # Only call callbacks on exception
            return super().__exit__(exc_type, exc_val, exc_tb)
        self._exit_callbacks.clear()

def rollback(msg):
    print("\tROLLBACK:", msg)

def callback(msg):
    print("\tCALLBACK:", msg)

print("Without Error")
with (
    contextlib.ExitStack() as always_stack,
    ExceptionExitStack() as exception_stack,
):
    always_stack.callback(callback, "via ExitStack")
    exception_stack.callback(rollback, "via ExceptionExitStack")


print("With Error")
with (
    contextlib.ExitStack() as always_stack,
    ExceptionExitStack() as exception_stack,
):
    always_stack.callback(callback, "via ExitStack")
    exception_stack.callback(rollback, "via ExceptionExitStack")
    raise RuntimeError("Application error")

Without Error
	CALLBACK: via ExitStack
With Error
	ROLLBACK: via ExceptionExitStack
	CALLBACK: via ExitStack

I was just playing around, but unless the inheritance totally borks an edge case, this is actually the exact implementation I'd go with. If you really despise it I'm not 100% closed to the context manager attack, but you'll need to feel pretty strongly or have a good failure case because I love how simple this implementation is 😂

One extension I can see is modifying the init to take the types of exception on which to trigger:

class ExceptionExitStack(contextlib.ExitStack):
    def __init__(self, *exceptions: type[Exception]):
        super().__init__()
        self._exception_types = None if len(exceptions) == 0 else exceptions

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type and (self._exception_types is None or exc_type in self._exception_types):
            return super().__exit__(exc_type, exc_val, exc_tb)
        self._exit_callbacks.clear()

def rollback(msg):
    print("\tROLLBACK:", msg)

print("With No Specifiers")
try:
    with ExceptionExitStack() as stack:
        stack.callback(rollback, "with no types")
        raise RuntimeError("Application error")
except Exception:
    pass

print("With Matching Specifier")
try:
    with ExceptionExitStack(RuntimeError) as stack:
        stack.callback(rollback, "with matching type")
        raise RuntimeError("Application error")
except Exception:
    pass

print("With Mis-Matching Specifier(s)")
try:
    with ExceptionExitStack(TypeError, ValueError) as stack:
        stack.callback(rollback, "with mis-matching types")
        raise RuntimeError("Application error")
except Exception:
    pass

With No Specifiers
	ROLLBACK: with no types
With Matching Specifier
	ROLLBACK: with matching type
With Mis-Matching Specifier(s)

[pmrv]

I don't feel strongly about this at all. If we're starting a decorator submodule, I feel a version of onerror would also be nice to offer options. I browsed the docs as well and it has some examples on how to do ExitStack subclasses. For your implementation I would only change to use pop_all rather than reaching into the private attribute and do a proper instance check to catch error subtypes

all(isinstance(exc_val, e) for e in self._exception_types)

Otherwise lgtm.

[pmrv]

Using ExistStack.push could be another option, but it would require that the rollback functions accept the exception values (or are taught to via a decorator).