Visual representation of sorts

sorts.py

#!/usr/bin/env python
"""Visual demonstration of sort algorithms"""

import curses
import random
import sys
import time


class SortBase:

    name = 'Just a base'

    def __init__(self, stdscr, speed):
        possible_delays = [1.0, 0.1, 0.01, 0.001, 0.0001]
        try:
            self.delay = possible_delays[int(speed) - 1]
        except (IndexError, ValueError):
            self.delay = possible_delays[2]
        curses.init_pair(1, curses.COLOR_GREEN, curses.COLOR_GREEN)
        curses.init_pair(2, curses.COLOR_BLUE, curses.COLOR_BLUE)
        self.stdscr = stdscr
        self.sort_pass = 0
        self.swap_count = 0
        num_items = int(curses.COLS / 2)
        max_value = curses.LINES - 5
        self.values = [random.randint(1, max_value) for _ in range(num_items)]
        stdscr.addstr(curses.LINES - 3, int(curses.COLS / 2 - len(self.name) / 2), self.name)

    def draw_all_values(self):
        self.draw_sort_pass()
        self.draw_swap_count()
        for idx, value in enumerate(self.values):
            self.erase_value(idx)
            self.draw_value(idx)
            self.stdscr.refresh()

    def draw_value(self, index):
        self.erase_value(index)
        for y in range(self.values[index]):
            self.stdscr.addstr(curses.LINES - 4 - y, index * 2, ' ', curses.A_REVERSE)

    def highlight_value(self, index, color_index=1):
        self.erase_value(index)
        for y in range(self.values[index]):
            self.stdscr.addstr(curses.LINES - 4 - y, index * 2, ' ', curses.color_pair(color_index))

    def erase_value(self, index):
        for y in range(curses.LINES - 5):
            self.stdscr.addstr(curses.LINES - 4 - y, index * 2, '_')

    def refresh(self):
        self.stdscr.refresh()
        time.sleep(self.delay)

    def swap_values(self, index_1, index_2):
        self.swap_count += 1
        self.highlight_value(index_1, 1)
        self.highlight_value(index_2, 2)
        self.refresh()
        self.values[index_1], self.values[index_2] = self.values[index_2], self.values[index_1]
        self.highlight_value(index_1, 2)
        self.highlight_value(index_2, 1)
        self.refresh()
        self.draw_value(index_1)
        self.draw_value(index_2)
        self.draw_sort_pass()
        self.draw_swap_count()
        self.refresh()

    def draw_sort_pass(self):
        self.stdscr.addstr(curses.LINES - 3, 0, f"Pass: {self.sort_pass+1}")

    def draw_swap_count(self):
        msg = f"Swap count: {self.swap_count}"
        left = curses.COLS - len(msg)
        self.stdscr.addstr(curses.LINES - 3, left, msg)

    def sort(self): # Overriden in the various sort implementations.
        self.values.sort()
        self.draw_values()

    def interactive_sort(self):
        self.draw_all_values()
        self.stdscr.addstr(0, 0, 'Press any key to begin')
        self.stdscr.getkey()
        self.stdscr.addstr(0, 0, '                      ')
        self.sort()
        self.stdscr.addstr(0, 0, 'Done - press any key to quit')
        self.stdscr.getkey()


class BubbleSort(SortBase):

    name = 'Bubble (simple) sort'

    def sort(self):
        values = self.values
        done = False
        while not done:
            done = True
            self.draw_sort_pass()
            for index in range(len(values)-1 - self.sort_pass):
                if values[index] > values[index+1]:
                    done = False
                    self.swap_values(index, index+1)
            self.sort_pass += 1


class SelectionSort(SortBase):

    name = 'Selection (simple) sort'

    def sort(self):
        values = self.values
        for sort_pass in range(len(values)):
            self.sort_pass = sort_pass
            max_index = 0
            for index in range(1, len(values) - sort_pass):
                if values[index] > values[max_index]:
                    max_index = index
            swap_index = len(values) - sort_pass - 1
            self.swap_values(max_index, swap_index)


def main(stdscr, sort_class, speed):
    curses.curs_set(0)          # Hide the cursor
    sort = sort_class(stdscr, speed)
    sort.interactive_sort()


if __name__ == '__main__':
    sorts = {'bubble': BubbleSort, 'selection': SelectionSort}
    chosen_sort = '' if len(sys.argv) < 2 else sys.argv[1].lower()
    speed = '3' if len(sys.argv) < 3 else sys.argv[2]
    if chosen_sort not in sorts:
        print(f"Usage: {sys.argv[0]} {'|'.join(sorts.keys())} 1-5")
    else:
        try:
            curses.wrapper(main, sorts[chosen_sort], speed)
        except KeyboardInterrupt:
            print('User quit.')