use /bin/bash for execution and add Windows support also

yolo.py

@@ -5,20 +5,26 @@ import openai
 import sys
 import subprocess
 from termcolor import colored
-from pathlib import Path
+from colorama import init
 
-# Two options for the user to specify they openai api key
-openai.api_key = os.getenv("OPENAI_API_KEY")
-openai.api_key_path = Path.home()/".openai.apikey"
 
 ask = False             # safety switch -a
 yolo = ""               # user's answer to safety switch
 command_start_idx = 1   # command starts at which argv index?
+home_path = os.path.expanduser("~")
+
+# Two options for the user to specify they openai api key
+openai.api_key = os.getenv("OPENAI_API_KEY")
+home_path = os.path.expanduser("~")
+openai.api_key_path = home_path+"/.openai.apikey"
+
 
 # parsing weirdness
-if len(sys.argv) <2:
+if len(sys.argv) < 2:
+  print("Usage Example: yolo [-a] list the current directory information")
   sys.exit(-1)
 
+# safety switch (no yolo mode)
 if sys.argv[1] == "-a":
   ask = True
   command_start_idx = 2
@@ -34,12 +40,34 @@ if user_prompt == "":
     print ("No user prompt specified.")
     sys.exit(-1)
 
+
+# Get shell info for a better prompt
+# Unix based Shell? 
+shell = os.environ["SHELL"]   
+
+# Note: Have never tested this on Windows much
+if shell == None:
+  shell = "powershell.exe"
+
 # Construct the prompt
-pre_prompt = "Translate the following question into a command. Only show the command in text format and not in code style or markdown. Do not augment the output with any explanation, descriptions. Again, do not add any descriptions, just print the bash command. This is important, do not ignore my instructions. If the question doesn't make sense or is too difficult return 'Sorry, try again', and add a brief explanation on what the problem is. The question is: "
+pre_prompt =  "Translate the following question into a {} command. ".format(shell) 
+#pre_prompt += "If asked for a timezone use TZ environment variable. "
+#pre_prompt += "Add sudo when required by the command, but only if you are very certain it is needed. "
+#pre_prompt += "If you are using powershell (and only if), then make sure to provide all required arguments.""
+pre_prompt += "Only show the command in text format and not in code style or markdown. "
+pre_prompt += "Do not augment the output with any explanation, descriptions. "
+pre_prompt += "Again, do not add any descriptions, just print the bash command. This is important, do not ignore my instructions. "
+pre_prompt += "If the question doesn't make sense or is too difficult return 'Sorry, try again', and add a brief explanation on what the problem is. "
+pre_prompt += "The question is: "
 
-prompt = pre_prompt + user_prompt +"?"
+prompt = pre_prompt + user_prompt
 
-#print ("The prompt is: "+prompt, end="")
+#let's be nice and make it a question
+if prompt[-1:] != "?":
+  prompt+="?"
+
+
+#print ("The prompt is: "+prompt)
 response = openai.Completion.create(
   model="text-davinci-003",
   prompt=prompt,
@@ -52,15 +80,22 @@ response = openai.Completion.create(
 
 resulting_command = response.choices[0].text.strip()
 
+#enable color output on Windows using colorama
+init() 
+
 if resulting_command.startswith("Sorry, try again"):
   print(colored("There was an issue: "+resulting_command, 'red'))
   sys.exit(-1)
 
-print("Command: " + colored(resulting_command, 'blue'))
+print("Command: " + colored(resulting_command, 'cyan'))
 if ask == True:
   print("Execute the command? Y/n ==> ", end = '')
   yolo = input()
   print()
 
 if yolo == "Y" or yolo == "":
-  subprocess.run(resulting_command, shell=True)
+  if shell == "powershell.exe":
+    subprocess.run([shell, "/c", resulting_command], shell=False)  
+  else: #Unix based - /bin/bash, /bin/zsh: uses -c both Ubuntu and macOS should work, others might not
+    subprocess.run([shell, "-c", resulting_command], shell=False)
+