Argument parsing via argparse

yolo.py

@@ -7,14 +7,16 @@ Sources:
 
 import os
 import platform
-import sys
 import subprocess
+import sys
 
-import openai
-import dotenv
+
+import argparse
 import distro
-import yaml
+import dotenv
+import openai
 import pyperclip
+import yaml
 
 from termcolor import colored
 from colorama import init
@@ -96,16 +98,10 @@ def get_full_prompt(user_prompt, shell):
 
     return prompt
 
-def print_usage():
+def print_config(config):
     """
-    Print usage information.
+    Print config information.
     """
-    print("Yolo v0.2.1 - by @wunderwuzzi23")
-    print()
-    print("Usage: yolo [-a] list the current directory information")
-    print("Argument: -a: Prompt the user before running the command (only useful when safety is off)")
-    print()
-
     print("Current configuration per yolo.yaml:")
     print("* Model        : " + str(config["model"]))
     print("* Temperature  : " + str(config["temperature"]))
@@ -230,38 +226,32 @@ def main():
     """
     Defined starting point of source code.
     """
+    parser = argparse.ArgumentParser(
+        description='AI bot that translates your question to a command.'
+    )
+    parser.add_argument('text', nargs='+', 
+                        help='A sequence of strings')
+    parser.add_argument("-s", "--safety", action='store_true', 
+                        help='Enable safety mode (only useful when safety is off)')
+    parser.add_argument("-c", "--config", action='store_true', 
+                        help='Print current configuration')
+    args = parser.parse_args()
+
+    user_prompt = " ".join(args.text)
+    ask_flag = args.safety
+    
     config = read_yaml_config()
     set_openai_api_key(config)
 
-    # set_api_key(config)
+    if (args.config):
+        print_config(config)
 
     # Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
     shell = os.environ.get("SHELL", "powershell.exe")
 
-    command_start_idx  = 1     # Question starts at which argv index?
-    ask_flag = False           # safety switch -a command line argument
-    yolo = ""                  # user's answer to safety switch (-a) question y/n
-
-    # Parse arguments and make sure we have at least a single word
-    if len(sys.argv) < 2:
-        print_usage()
-        sys.exit(-1)
-
-    # Safety switch via argument -a (local override of global setting)
-    # Force Y/n questions before running the command
-    if sys.argv[1] == "-a":
-        ask_flag = True
-        command_start_idx = 2
-
-    # To allow easy/natural use we don't require the input to be a
-    # single string. So, the user can just type yolo what is my name?
-    # without having to put the question between ''
-    arguments = sys.argv[command_start_idx:]
-    user_prompt = " ".join(arguments)
-    
     #Enable color output on Windows using colorama
     init()
-    
+
     res_command = call_open_ai(config, shell, user_prompt)
     check_for_issue(res_command)
     check_for_markdown(res_command)
@@ -270,4 +260,4 @@ def main():
     evaluate_input(user_input, res_command)
 
 if __name__ == "__main__":
-    main()
+    main()