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

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?"

In a notebook. thing_instance.do_stuff(). And actually I use a decorator to do the state toggling. So `do_stuff` is decorated. But `do_other_stuff` is not. — Alexander Soare, Sep 18 '20 at 10:56
Does this answer your question? [Doing something before program exit](https://stackoverflow.com/questions/3850261/doing-something-before-program-exit) — Georgy, Sep 25 '20 at 07:07

score 5 · Accepted Answer · answered Sep 18 '20 at 11:03

5

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.

answered Sep 18 '20 at 11:03

Gilles 'SO- stop being evil'

104,111
38
209
254

All good answers here but this was the most concise. I ended up using `except Exception as e` so that I could then go ahead and raise it after cleaning up. – Alexander Soare Sep 18 '20 at 11:16
3

@AlexanderSoare `finally` re-raises the exception. It's uncommon to need to catch an exception and re-raise it, mainly it's done when you want to log the exception. – Gilles 'SO- stop being evil' Sep 18 '20 at 11:19

TayTay · Answer 2 · 2020-09-18T11:03:33.470

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

MisterMiyagi · Answer 3 · 2020-09-18T11:09:00.623

Stopping execution is equivalent to raising KeyboardInterrupt. You can handle this exception like any other.

Use a try:/finally: to always perform cleanup:

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

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

If a cleanup pattern is needed often, a context manager can be used to perform pre-defined enter and cleanup actions. For example, if some_state is to prevent execution of other code, a Lock is appropriate.

import threading

class Thing():
    def __init__(self):
        self.some_lock = threading.Lock()

    def do_stuff(self):
        with self.some_lock:  # prevent doing other stuff while in this block
            # do stuff which may take some time - and user may quit here

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

3 Answers3