Make everything more pythonic & cleaner

pridefetch

@@ -1,109 +1,123 @@
 #!/usr/bin/env python3
 
-import argparse
+from argparse import ArgumentParser
-import distro
+from distro import name as distro_name
-import random
+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 time import strftime
+from datetime import timedelta
-from time import gmtime
-from subprocess import run
 
 # 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],
+    "straight": [0, 255, 0, 255, 0],
-    'gay': [196, 208, 226, 28, 20, 90],
+    "gay": [196, 208, 226, 28, 20, 90],
-    'bisexual': [198, 198, 97, 25, 25],
+    "bisexual": [198, 198, 97, 25, 25],
-    'lesbian': [202, 209, 255, 255, 168, 161],
+    "lesbian": [202, 209, 255, 255, 168, 161],
-    'pansexual': [198, 198, 220, 220, 39, 39],
+    "pansexual": [198, 198, 220, 220, 39, 39],
-    'trans': [81, 211, 255, 211, 81],
+    "trans": [81, 211, 255, 211, 81],
-    'nonbinary': [226, 226, 255, 255, 98, 98, 237, 237],
+    "nonbinary": [226, 226, 255, 255, 98, 98, 237, 237],
-    'demiboy': [244, 249, 117, 255, 117, 249, 244],
+    "demiboy": [244, 249, 117, 255, 117, 249, 244],
-    'demigirl': [244, 249, 218, 255, 218, 249, 244],
+    "demigirl": [244, 249, 218, 255, 218, 249, 244],
-    'genderfluid': [211, 255, 128, 0, 63],
+    "genderfluid": [211, 255, 128, 0, 63],
-    'aromantic': [71, 149, 255, 249, 0],
+    "aromantic": [71, 149, 255, 249, 0],
-    'agender': [0, 251, 255, 149, 255, 251, 0],
+    "agender": [0, 251, 255, 149, 255, 251, 0],
-    'asexual': [0, 0, 242, 242, 255, 255, 54, 54],
+    "asexual": [0, 0, 242, 242, 255, 255, 54, 54],
-    'graysexual': [54, 242, 255, 242, 54],
+    "graysexual": [54, 242, 255, 242, 54],
 }
 
-def color256(col, bg_fg):
+# Each box for each part of the flag
-    # Hacky alias around manually typing out escape codes every time 
+box = " "
-    return f'\033[{48 if bg_fg == "bg" else 38};5;{col}m'
 
-def draw_info(flag_name):
+# When printed, reset will end the color of the row
-    box = ' '
+reset = "\033[0m\033[39m"
-    width = 20
 
+
+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]
-    curr_row = 0
+    width = width or 20
-
-    # Store the output of uname -srm
-    uname_info = run(['uname', '-srm'], capture_output=True)
-    uname_info = uname_info.stdout.decode().strip()
 
     # 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")
 
-    reset = '\033[0m\033[39m'
+    # Print a blank line to separate the flag from the terminal prompt
-
     print()
-    for row in flag:
+
+    for row_number, row in enumerate(flag):
         # Alternate displaying the information based on the current row
-        if curr_row == 0:
+        match row_number:
-            color = color256(flag[curr_row], "fg")
+            case 0:
+                # Get the username and hostname
+                user = getuser()
+                host = gethostname()
 
-            user = getuser()
+                # Username and hostname
-            host = gethostname()
+                color = color256(flag[row_number], "fg")
-            row_info = f'{color}\033[1m{user}@{host}{reset}'
+                row_info = f"{color}\033[1m{user}@{host}{reset}"
-        elif curr_row == 1:
-            distribution = distro.name()
-            row_info = f'{row_color}os      {reset}{distribution or "N/A"}'
-        elif curr_row == 2:
-            arch = uname_info.split(' ')[2]
-            row_info = f'{row_color}arch    {reset}{arch}'
-        elif curr_row == 3:
-            kernel = uname_info.split(' ')[1]
-            row_info = f'{row_color}kern    {reset}{kernel}'
-        elif curr_row == 4:
-            uptime = run(['uptime'], capture_output=True)
-            time = uptime.stdout.decode().split(" ")[1]
 
-            row_info = f'{row_color}uptime  {reset}{time}'
+            case 1:
-        else:
+                # Distribution
-            row_info = ''
+                distribution = distro_name()
-        
+                row_info = f"{row_color}os      {reset}{distribution or 'N/A'}"
-        if row != 'P':
-            print(f' {color256(row, "bg")}{box * width}\033[49m {row_info}')
-        else:
-            # Print just the info, along with the padding for the box
-            print(f' {box * width}\033[49m {row_info}')
-            
-        curr_row += 1
 
+            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():
-    parser = argparse.ArgumentParser()
+    # Argument configuration
-    parser.add_argument('-f', '--flag',     help='Displays the chosen FLAG')
+    parser = ArgumentParser()
-    parser.add_argument('-c', '--choose',   help='Choose a flag at random from this list')
+    parser.add_argument("-f", "--flag", help="displays the chosen flag")
-    parser.add_argument('-l', '--list',     action='store_true', help='Lists all the flags')
+    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_info(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(',')
+        flag_choices = args.choose.split(",")
-        draw_info(random.choice(flag_choices))
+        draw_fetch(random_choice(flag_choices))
+
     if args.list:
         # List out all the available flags
-        print('Available flags:')
+        print(f"Available flags:\n{', '.join(flags)}")
-        print(', '.join(flags))
-    
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()