PrideFetch/pridefetch

#!/usr/bin/env python3

from argparse import ArgumentParser
from distro import name as distro_name
from platform import machine, release
from random import choice as random_choice
from time import clock_gettime, CLOCK_BOOTTIME
from getpass import getuser
from socket import gethostname
from datetime import timedelta

# Define a dictionary of all the flags and their colors
# Each color is the color for an individual row in the flag
flags = {
    "straight": [0, 255, 0, 255, 0],
    "gay": [196, 208, 226, 28, 20, 90],
    "bisexual": [198, 198, 97, 25, 25],
    "lesbian": [202, 209, 255, 255, 168, 161],
    "pansexual": [198, 198, 220, 220, 39, 39],
    "trans": [81, 211, 255, 211, 81],
    "nonbinary": [226, 226, 255, 255, 98, 98, 237, 237],
    "demiboy": [244, 249, 117, 255, 117, 249, 244],
    "demigirl": [244, 249, 218, 255, 218, 249, 244],
    "genderfluid": [211, 255, 128, 0, 63],
    "aromantic": [71, 149, 255, 249, 0],
    "agender": [0, 251, 255, 149, 255, 251, 0],
    "asexual": [0, 0, 242, 242, 255, 255, 54, 54],
    "graysexual": [54, 242, 255, 242, 54],
}

# Each box for each part of the flag
box = " "

# When printed, reset will end the color of the row
reset = "\033[0m\033[39m"


def color256(col: int, bg_fg: str) -> str:
    # Hacky alias around manually typing out escape codes every time
    return f"\033[{48 if bg_fg == 'bg' else 38};5;{col}m"


def draw_fetch(flag_name: str, width: int = None):
    # Load configuration variables & flag data
    flag = flags[flag_name]
    width = width or 20

    # Make sure that the row color is different to the color of the hostname
    row_color = color256(flag[1] if flag[0] != flag[1] else flag[2], "fg")

    # Print a blank line to separate the flag from the terminal prompt
    print()

    for row_number, row in enumerate(flag):
        # Alternate displaying the information based on the current row
        match row_number:
            case 0:
                # Get the username and hostname
                user = getuser()
                host = gethostname()

                # Username and hostname
                color = color256(flag[row_number], "fg")
                row_info = f"{color}\033[1m{user}@{host}{reset}"

            case 1:
                # Distribution
                distribution = distro_name()
                row_info = f"{row_color}os      {reset}{distribution or 'N/A'}"

            case 2:
                # Architecture
                arch = machine()
                row_info = f"{row_color}arch    {reset}{arch}"

            case 3:
                # Kernel version
                kernel = release()
                row_info = f"{row_color}kern    {reset}{kernel}"

            case 4:
                # Uptime
                time = str(timedelta(seconds=clock_gettime(CLOCK_BOOTTIME))).split(".", 1)[0]
                row_info = f"{row_color}uptime  {reset}{time}"

            case _:
                # Empty row
                row_info = ""

        # Print the row next to the flag
        print(f" {color256(row, 'bg')}{box * width}\033[49m {row_info}")

    # Print a blank line to separate the flag from the terminal prompt
    print()


def main():
    # Argument configuration
    parser = ArgumentParser()
    parser.add_argument("-f", "--flag", help="displays the chosen flag")
    parser.add_argument("-c", "--choose", help="choose a flag at random from a list seperated by commas")
    parser.add_argument("-w", "--width", help="choose a custom width for the flag", type=int)
    parser.add_argument("-l", "--list", help="lists all the flags that can be displayed", action="store_true")

    # Parse the arguments
    args = parser.parse_args()

    if args.flag:
        # Draw the flag
        draw_fetch(args.flag, args.width)

    if args.choose:
        # Choose a flag at random from a list of comma-seperated flags
        flag_choices = args.choose.split(",")
        draw_fetch(random_choice(flag_choices))

    if args.list:
        # List out all the available flags
        print(f"Available flags:\n{', '.join(flags)}")


if __name__ == "__main__":
    main()