terminal-rain/colorful-rain.py

#!/usr/bin/env python3

'''
terminal-rain - colorful-rain.py  
Copyright (c) 2021 Julian Müller (ChaoticByte)
License: MIT  
'''

import os, random, sys, time

if len(sys.argv) > 1 and sys.argv[1] == "b":
    colorpalette = [
        "\u001b[38;5;198m",
        "\u001b[38;5;198m",
        "\u001b[38;5;99m",
        "\u001b[38;5;33m",
        "\u001b[38;5;33m"
    ]
else:
    colorpalette = [
        "\u001b[38;5;196m",
        "\u001b[38;5;202m",
        "\u001b[38;5;220m",
        "\u001b[38;5;82m",
        "\u001b[38;5;39m",
        "\u001b[38;5;129m"
    ]

resetcolor = "\u001b[0m" 

class Raindrop():
    def __init__(self, posX, color):
        self.posX = posX
        self.posY = 0
        self.symbol = "|"
        self.color = color

# array for the raindrops
raindrops = []

def generate_new_raindrops(n):
    # generate n new raindrops per frame
    for _ in range(n):
        raindrops.append(Raindrop(random.randint(1, 1000) % (width + 1), colorpalette[random.randint(0, len(colorpalette) - 1)]))

def remove_old_raindrops(x, n):
    # when there are more than x raindrops, delete the last n
    if len(raindrops) > x:
        for _ in range(n):
            raindrops.pop(0)

def calculate_raindrop_locations():
    # calculate the new position for every raindrop
    for elem_ in raindrops:
        if elem_.posY <= height - 1:
            elem_.posY += 1
        else:
            elem_.symbol = "_"

def write():
    # clear the screen
    sys.stdout.write("\u001b[2J")
    # write the raindrops to the terminal
    for elem_ in raindrops:
        sys.stdout.write(f"\u001b[{elem_.posY};{elem_.posX}f")
        sys.stdout.write(elem_.color)
        sys.stdout.write(elem_.symbol)
        sys.stdout.write(resetcolor)
    sys.stdout.write(f"\u001b[H")
    # flush the changes to the terminal
    sys.stdout.flush()


while True:
    width, height = os.get_terminal_size()
    generate_new_raindrops(3)
    remove_old_raindrops(230, 3)
    calculate_raindrop_locations()
    write()
    time.sleep(0.032)