Simple Matrix rain effect for the linux terminal

matrix-effect.py

import fcntl
import time
import random
import struct
import signal
import sys
import termios
import argparse

DEFAULT_SPEED = 0.05
DEFAULT_COLOR = 'green'

COLORS = {
    'gray': 30,
    'grey': 30,
    'red': 31,
    'green': 32,
    'yellow': 33,
    'blue': 34,
    'purple': 35,
    'lightblue': 36,
    'white': 37
}

def translate_text_color(color_name):
    return COLORS.get(color_name) if color_name in COLORS.keys() else COLORS.get(DEFAULT_COLOR)

class MatrixRain:
    def __init__(self, text_color, speed):
        self.items = [unichr(i) for i in range(0x30a1, 0x30ff + 1)]
        self.screen_height, self.screen_width = struct.unpack('hh', fcntl.ioctl(1, termios.TIOCGWINSZ, '1234'))
        self.row_states = [0 for i in range(0, int(self.screen_width / 2))]
        self.row_grid = []
        self.rain_text_color = text_color if text_color else COLORS.get(DEFAULT_COLOR)
        self.rain_speed = speed if speed else DEFAULT_SPEED

    def create_character(self, single_character, index):
        if index >= 30:
            text_color = 39
        else:
            text_color = self.rain_text_color

        return '\033[1m\033[%dm%s\033[0m' % (text_color, single_character)

    def create_row(self):
        output = ''
        for x in range(0, len(self.row_states)):
            old_state = self.row_states[x]
            if self.row_states[x] == 0:
                show_character = random.randint(0, 10)
                if show_character > 4:
                    self.row_states[x] = random.randint(5, 40)
                else:
                    self.row_states[x] = random.randint(-40, -1)
            elif self.row_states[x] > 0:
                self.row_states[x] -= 1
            elif self.row_states[x] < 0:
                self.row_states[x] += 1

            if self.row_states[x] > 0:
                if old_state == 0:
                    output += self.create_character(self.items[random.randint(0, len(self.items) - 1)], 30)
                else:
                    output += self.create_character(self.items[random.randint(0, len(self.items) - 1)], self.row_states[x])
            else:
                output += '  '

        return output

    def create_empty_row(self):
        output = ''
        for x in range(0, len(self.row_states)):
            output += '  '
        return output

    def print_missing_rows(self, existing_rows = 0):
        if existing_rows < self.screen_height:
            for x in range(0, self.screen_height - existing_rows):
                sys.stdout.write(self.create_empty_row())
                sys.stdout.flush()

    def run_matrix(self):
        if len(self.row_grid) == self.screen_height:
            self.row_grid.pop(len(self.row_grid) - 1)
            self.row_grid.insert(0, self.create_row())
        else:
            self.row_grid.insert(0, self.create_row())

        for x in range(0, len(self.row_grid)):
            sys.stdout.write(self.row_grid[x])
            sys.stdout.flush()

        self.print_missing_rows(len(self.row_grid))

        time.sleep(self.rain_speed)

def signal_handler(signal, frame):
    print('you left the matrix')
    sys.exit(0)

def main():
    signal.signal(signal.SIGINT, signal_handler)

    parser = argparse.ArgumentParser(description='Matrix rain effect in terminal.')
    parser.add_argument('-c', '--color', help='color of the characters', type=str)
    parser.add_argument('-s', '--speed', help='speed of the rain', type=float)

    args = parser.parse_args()

    rain_text_color = translate_text_color(args.color) if args.color else None
    rain_speed = args.speed if args.speed else None

    matrix_rain = MatrixRain(rain_text_color, rain_speed)
    while True:
        matrix_rain.run_matrix()

    sys.exit(0)

if __name__ == "__main__":
    main()