Search code examples
pythonexceptionspinnerkeyboardinterrupt

Better keyboard interrupt detection for this threaded Spinner class


Ok, I've wrote this class based in a bunch of others Spinner classes that I've googled in Google Code Search.

It's working as intended, but I'm looking for a better way to handle KeyboardInterrupt and SystemExit exceptions. Is there better approaches?

Here's my code:

#!/usr/bin/env python
import itertools
import sys
import threading

class Spinner(threading.Thread):
    '''Represent a random work indicator, handled in a separate thread'''

    # Spinner glyphs
    glyphs = ('|', '/', '-', '\\', '|', '/', '-')
    # Output string format
    output_format = '%-78s%-2s'
    # Message to output while spin
    spin_message = ''
    # Message to output when done
    done_message = ''
    # Time between spins
    spin_delay = 0.1

    def __init__(self, *args, **kwargs):
        '''Spinner constructor'''
        threading.Thread.__init__(self, *args, **kwargs)
        self.daemon = True
        self.__started = False
        self.__stopped = False
        self.__glyphs = itertools.cycle(iter(self.glyphs))

    def __call__(self, func, *args, **kwargs):
        '''Convenient way to run a routine with a spinner'''
        self.init()
        skipped = False

        try:
            return func(*args, **kwargs)
        except (KeyboardInterrupt, SystemExit):
            skipped = True
        finally:
            self.stop(skipped)

    def init(self):
        '''Shows a spinner'''
        self.__started = True
        self.start()

    def run(self):
        '''Spins the spinner while do some task'''
        while not self.__stopped:
            self.spin()

    def spin(self):
        '''Spins the spinner'''
        if not self.__started:
            raise NotStarted('You must call init() first before using spin()')

        if sys.stdin.isatty():
            sys.stdout.write('\r')
            sys.stdout.write(self.output_format % (self.spin_message,
                                                   self.__glyphs.next()))
            sys.stdout.flush()
            time.sleep(self.spin_delay)

    def stop(self, skipped=None):
        '''Stops the spinner'''
        if not self.__started:
            raise NotStarted('You must call init() first before using stop()')

        self.__stopped = True
        self.__started = False

        if sys.stdin.isatty() and not skipped:
            sys.stdout.write('\b%s%s\n' % ('\b' * len(self.done_message),
                                           self.done_message))
            sys.stdout.flush()

class NotStarted(Exception):
    '''Spinner not started exception'''
    pass

if __name__ == '__main__':
    import time

    # Normal example
    spinner1 = Spinner()
    spinner1.spin_message = 'Scanning...'
    spinner1.done_message = 'DONE'
    spinner1.init()
    skipped = False

    try:
        time.sleep(5)
    except (KeyboardInterrupt, SystemExit):
        skipped = True
    finally:
        spinner1.stop(skipped)

    # Callable example
    spinner2 = Spinner()
    spinner2.spin_message = 'Scanning...'
    spinner2.done_message = 'DONE'
    spinner2(time.sleep, 5)

Thank you in advance.


Solution

  • You probably don't need to worry about catching SystemExit as it is raised by sys.exit(). You might want to catch it to clean up some resources just before your program exits.

    The other way to catch KeyboardInterrupt is to register a signal handler to catch SIGINT. However for your example using try..except makes more sense, so you're on the right track.

    A few minor suggestions:

    • Perhaps rename the __call__ method to start, to make it more clear you're starting a job.
    • You might also want to make the Spinner class reusable by attaching a new thread within the start method, rather than in the constructor.
    • Also consider what happens when the user hits CTRL-C for the current spinner job -- can the next job be started, or should the app just exit?
    • You could also make the spin_message the first argument to start to associate it with the task about to be run.

    For example, here is how someone might use Spinner:

    dbproc = MyDatabaseProc()
    spinner = Spinner()
    spinner.done_message = 'OK'
    try:
        spinner.start("Dropping the database", dbproc.drop, "mydb")
        spinner.start("Re-creating the database", dbproc.create, "mydb")
        spinner.start("Inserting data into tables", dbproc.populate)
        ...
    except (KeyboardInterrupt, SystemExit):
        # stop the currently executing job
        spinner.stop()
        # do some cleanup if needed..
        dbproc.cleanup()