Python - Is there a way to make a function clean up gracefully if the user tries to stop execution midway? [closed]

Question

Say I have a class like

class Thing():

    def __init__(self):
        self.some_state = False

    def do_stuff(self):
        self.some_state = True
        # do stuff which may take some time - and user may quit here
        self.some_state = False

    def do_other_stuff(self):
        # do stuff which depends on `some_state` being False

And I want to make sure that if a user executes this in a notebook by running:

thing = Thing()
thing.do_stuff()

then presses "stop execution" while running, some_state toggles back to False. That way do_other_stuff will work as intended. Is there a way to do some graceful clean up?

Note: Although my example is quite specific, my question is generally: "Can I do graceful cleanup?"

Answer 1

Stopping the execution raises the KeyboardInterrupt exception, so you need to handle that exception. But in fact, you probably need to reset some_state if the code exits with other exceptions as well. Even if you don't raise exceptions explicitly, they might occur due to a bug in your code or due to running out of memory. So do the cleanup in a finally clause.

    def do_stuff(self):
        self.some_state = True
        try:
            # do stuff which may take some time - and user may quit here
        finally:
            self.some_state = False

If many methods require the same cleanup, you can use a decorator as illustrated by Tgsmith61591.

Answer 2

Here's one way you can do this with a decorator:

def graceful_cleanup(f):
    def _inner(*args, **kwargs):
        self = args[0]
        try:
            return f(*args, **kwargs)
        except KeyboardInterrupt:
            self.some_state = False
    return _inner


class Thing():
    def __init__(self):
        self.some_state = False
        
    @graceful_cleanup
    def do_stuff(self):
        self.some_state = True
        time.sleep(5)
        self.some_state = False

If you try this out you'll see it resets its state:

In [2]: t = Thing()

In [3]: t.do_stuff()
^C
In [4]: t.some_state
Out[4]: False

Another cool way to do this is with a context manager:

import contextlib


@contextlib.contextmanager
def patch_attribute(instance, key, value):
    original_value = getattr(instance, key)
    try:
        setattr(instance, key, value)
        yield
    finally:
        setattr(instance, key, original_value)


class Thing():
    def __init__(self):
        self.some_state = False

    def do_stuff(self):
        with patch_attribute(self, "some_state", True):
            assert self.some_state is True
        assert not self.some_state  # show it resets outside the ctx