diff --git a/.gitignore b/.gitignore index ed8ebf5..341e31e 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ -__pycache__ \ No newline at end of file +__pycache__ +.idea \ No newline at end of file diff --git a/src/__main__.py b/src/__main__.py index 10707da..c8f0082 100755 --- a/src/__main__.py +++ b/src/__main__.py @@ -4,7 +4,7 @@ from argparse import ArgumentParser from datetime import timedelta from random import choice as random_choice -from time import clock_gettime, CLOCK_BOOTTIME +import color # Title - user@hostname from getpass import getuser @@ -13,7 +13,8 @@ from socket import gethostname # System info modules from platform import platform as system from platform import release as kernel -# from platform import machine as architecture +from time import clock_gettime, CLOCK_BOOTTIME +from platform import machine as architecture from distro import name as distribution from modules.packages import get_num_packages as packages @@ -34,38 +35,59 @@ flags = { "aromantic": [71, 149, 255, 249, 0], "agender": [0, 251, 255, 149, 255, 251, 0], "asexual": [0, 242, 255, 54], - "graysexual": [54, 242, 255, 242, 54], + "graysexual": [54, 242, 255, 242, 54] } -# When printed, reset will end the color of the row -reset = "\033[0m\033[39m" +# A dictionary of all the available stats +stats = { + "os": lambda: distribution() or system() or 'N/A', + "arch": lambda: architecture() or 'N/A', + "pkgs": lambda: packages() or 'N/A', + "kernel": lambda: kernel() or system() or 'N/A', + "uptime": lambda: str(timedelta(seconds=clock_gettime(CLOCK_BOOTTIME))).split('.', 1)[0] +} -def color256(col: int, bg_fg: str) -> str: - # Alias to avoid manually typing out escape codes every time - return f"\033[{48 if bg_fg == 'bg' else 38};5;{col}m" +def generate_fetch(flag_name: str, show_stats: list = None, width: int = None) -> (list, int, list): + """ + Generates variables needed for a fetch + :param flag_name: The name of the flag to use + :param show_stats: Stats to show in the fetch + :param width: Custom width of the flag + :return: Generated flag data + """ - -def draw_fetch(flag_name: str, width: int = None): # Load the chosen flag from the dictionary of flags flag = flags[flag_name] # 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") + row_color = color.color256(flag[1] if flag[0] != flag[1] else flag[2], "fg") - # The fetch data (system info) to be displayed - row_data = [ - f"{color256(flag[0], 'fg') if flag[0] != 0 else color256(242, 'fg')}" - f"\033[1m{getuser()}@{gethostname()}{reset}", - f"{row_color}os {reset}{distribution() or system() or 'N/A'}", - # f"{row_color}arch {reset}{architecture() or 'N/A'}", - f"{row_color}pkgs {reset}{packages() or 'N/A'}", - f"{row_color}kernel {reset}{kernel() or system() or 'N/A'}", - f"{row_color}uptime {reset}{str(timedelta(seconds=clock_gettime(CLOCK_BOOTTIME))).split('.', 1)[0]}" + # Set default stats to show in the fetch + show_stats = show_stats or ["os", "pkgs", "kernel", "uptime"] + + # Initialise the fetch data (system info) to be displayed with the user@hostname + data = [ + f"{color.color256(flag[0], 'fg') if flag[0] != 0 else color.color256(242, 'fg')}" + f"\033[1m{getuser()}@{gethostname()}{color.clear}", ] + # Add the chosen stats to the list row_data + for stat in show_stats: + # Calculate the value for the stat by running its function + value = stats[stat]() + + # Calculate the correct amount of spaces to keep the stat values in line with each other + spaces = ((len(max(show_stats)) - len(stat)) + 1) * " " + + # Generate a row with color, stat name and its value + row = f"{row_color}{stat}:{spaces}{color.clear}{value}" + + # Add the row to the data + data.append(row) + # Until the flag is a greater length than the data - while len(flag) < len(row_data): + while len(flag) < len(data): # If the data is greater than the flag length then duplicate the length of the flag flag = [element for element in flag for _ in (0, 1)] @@ -73,49 +95,186 @@ def draw_fetch(flag_name: str, width: int = None): width = width or round(len(flag) * 1.5 * 3) # Ensures nothing is printed for empty lines - row_data.append("") + data.append("") + + # Return all the flag information ready for drawing + return flag, width, data + + +def draw_fetch(flag: list, width: int, data: list) -> None: + """ + Draws a fetch to the screen + :param flag: The flag as a list of colors + :param width: Width of the flag rows + :param data: System stats data + """ # Print a blank line to separate the flag from the terminal prompt print() for index, row in enumerate(flag): # Print out each row of the fetch - print(f" {color256(row, 'bg')}{' ' * width}\033[49m{reset} {row_data[min(index, len(row_data) - 1)]}{reset}") + print(f" {color.color256(row, 'bg')}{' ' * width}\033[49m{color.clear} " # Flag rows + f"{data[min(index, len(data) - 1)]}{color.clear}") # Stats rows - # Print a blank line to separate the flag from the terminal prompt + # Print a blank line again to separate the flag from the terminal prompt print() +def create_fetch(flag_name: str, show_stats: list = None, width: int = None) -> None: + """ + Creates a fetch, by generating and then drawing it + :param flag_name: The name of the flag to use + :param show_stats: Stats to show in the fetch + :param width: Custom width of the flag + """ + + # Check if the flag exists in the dictionary of flags + assert flag_name in flags.keys(), f"flag '{flag_name}' is not a valid flag" + + # Generate a fetch with the given info + flag, width, data = generate_fetch(flag_name, show_stats, width) + + # Draw the fetch + draw_fetch(flag, width, data) + + +def check_valid_argument(arg_flag: str, argument: str, valid_arguments: list) -> bool: + """ + Checks if an argument is valid by checking if it's in a list of valid arguments + :param arg_flag: The argument flag e.g. --random, --stats etc. + :param argument: A user inputted argument + :param valid_arguments: The valid list of arguments to check against + :return: True if the argument is valid, False if not + """ + + # Check if argument is valid, by checking if it is not in valid_arguments + if argument not in valid_arguments: + _print_error(f"Invalid argument '{argument}' given for '{arg_flag}'", + f"must be one of '{', '.join(valid_arguments)}'") + return False + + else: + return True + + +def check_valid_arguments(arg_flag: str, arguments: list, valid_arguments: list) -> bool: + """ + Checks if arguments are valid by checking if they are in a list of valid arguments + :param arg_flag: The argument flag e.g. --random, --stats etc. + :param arguments: A list of user inputted arguments + :param valid_arguments: The valid list of arguments to check against + :return: True if the arguments are valid, False if not + """ + + # If there are any arguments remaining + if len(arguments) > 0: + for argument in arguments: + # If the argument isn't in valid_arguments, it isn't valid + if not check_valid_argument(arg_flag, argument, valid_arguments): + return False + + # Otherwise, the user must have typed comma(s) without any arguments + else: + _print_error(f"No arguments given for '{arg_flag}'", + f"must be one of '{', '.join(valid_arguments)}'") + return False + + return True + + +def parse_comma_arguments(arg_flag: str, comma_arguments: str, valid_arguments: list) -> list: + """ + Parses comma seperated arguments and checks if they are valid + :param arg_flag: The argument command line flag e.g. --random, --stats etc. + :param comma_arguments: Raw string of user inputted arguments including commas + :param valid_arguments: The valid list of arguments to check against + :return: Parsed arguments if valid, exits the program if invalid + """ + + # Separate arguments into a list + arguments = comma_arguments.split(",") + + # Remove whitespaces from the list of arguments + arguments = [argument.strip() for argument in arguments if argument.strip()] + + # Check if the passed arguments are valid, if not, exit with an error + if not check_valid_arguments(arg_flag, arguments, valid_arguments): + exit(1) + + # Otherwise return the arguments + else: + return arguments + + +def _print_error(error: str, help_message: str = None) -> None: + """ + Prints an error message with optionally an extra help message + :param error: Error message to print + :param help_message: Optional help message + :return: + """ + + # Print out the error message + print(f"{color.bold}{color.red}Error: {error}{color.clear}") + + # If the help message was given, print it out + if help_message: + print(f" {color.red}╰> {help_message}{color.clear}") + + def main(): - # Argument configuration - options + """ + Main function that evaluates command line arguments + """ + + # Argument configuration - pridefetch command line options parser = ArgumentParser() - parser.add_argument("-f", "--flag", help="displays the chosen flag") - parser.add_argument("-r", "--random", help="randomly choose a flag from a list seperated by commas") + parser.add_argument("-l", "--list", help="lists all flags and stats that can be displayed", action="store_true") + parser.add_argument("-a", "--all-stats", help="use all available stats (overrides '--stats')", action="store_true") + parser.add_argument("-f", "--flag", help="displays a flag of your choice") + parser.add_argument("-r", "--random", help="randomly choose a flag from a comma-seperated list") + parser.add_argument("-s", "--stats", help="choose the stats to appear from a comma-seperated list") 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 (collect) any arguments args = parser.parse_args() + if args.all_stats: + # Add all the available stats to show_stats + show_stats = list(stats) + + elif args.stats: + # Parse chosen statistics arguments if they exist + show_stats = parse_comma_arguments("--stats", args.stats, list(stats)) + + else: + # Otherwise, use the default stats + show_stats = None + if args.flag: - # Check if the flag exists in the dictionary of flags - assert args.flag in flags.keys(), f"flag '{args.flag}' is not a valid flag" + # Check if the flag is a valid flag + if not check_valid_argument("--flag", args.flag, list(flags)): + exit(1) # Draw the chosen flag and system information - draw_fetch(args.flag, args.width) + create_fetch(args.flag, show_stats, args.width) elif args.random: - # Choose a flag at random from a list of comma-seperated flags - flag_choices = args.random.split(",") - draw_fetch(random_choice(flag_choices), args.width) + # Parse chosen random flag arguments if they exist + flag_choices = parse_comma_arguments("--random", args.random, list(flags)) + + # Draw a randomly selected flag from the list + create_fetch(random_choice(flag_choices), show_stats, args.width) elif args.list: - # List out all the available flags - print(f"Available flags:\n{', '.join(flags)}") + # List out all the available flags and stats + print(f"{color.bold}Available flags:{color.clear}\n{', '.join(flags)}\n\n" + f"{color.bold}Available stats:{color.clear}\n{', '.join(stats)}") else: # By default, draw the classic flag - draw_fetch("classic", args.width) + create_fetch("classic", show_stats, args.width) if __name__ == "__main__": diff --git a/src/color.py b/src/color.py new file mode 100644 index 0000000..36f811c --- /dev/null +++ b/src/color.py @@ -0,0 +1,10 @@ +# ASCII color codes for text +clear = "\033[0m\033[39m" +bold = "\033[1m" +red = "\033[31m" + + +def color256(col: int, bg_fg: str) -> str: + # Alias to avoid manually typing out escape codes every time for flags + return f"\033[{48 if bg_fg == 'bg' else 38};5;{col}m" + diff --git a/src/modules/packages.py b/src/modules/packages.py index 0fd6754..47f7ae1 100644 --- a/src/modules/packages.py +++ b/src/modules/packages.py @@ -23,6 +23,7 @@ def get_num_packages() -> (int, bool): try: # Get the length of the output of the command - the number of packages return len(check_output(command.split(" ")).decode("utf-8").split("\n")) - 1 + except FileNotFoundError: # If the command doesn't exist, skip it pass