Fix most of the big bugs finally, v0.2.7 release

Author: morganpartee

codegpt/gpt_interface.py

@@ -8,6 +8,7 @@
 try:
     nltk.data.find('tokenizers/punkt')
 except LookupError:
+    typer.secho("Downloading punkt for nltk... Only once!", fg=typer.colors.GREEN, bold=True)
     nltk.download('punkt', quiet=True)
 
 def confirm_send(prompt, max_tokens=4000, yes=False):
@@ -27,7 +28,7 @@ def confirm_send(prompt, max_tokens=4000, yes=False):
 
     return max_tokens
 
-def send_iffy_edit(prompt: str, code: Dict[str, str], clipboard: bool = False) -> str:
+def send_iffy_edit(prompt: str, code: Dict[str, str], clipboard: bool = False, yes=False) -> str:
     full_prompt = "You are an expert developer. Given the following code:\n\n"
 
     for filename, code_string in code.items():
@@ -39,19 +40,18 @@ def send_iffy_edit(prompt: str, code: Dict[str, str], clipboard: bool = False) -
 
     if not clipboard:
         full_prompt += dedent("""
+        You MUST output complete files.
         
-        Answer in this exact format:
+        For each file you wish to output (only if you modified it or made a new one),
+        answer in this exact format:
         
         filename:
         > <the filename to be output>
         explanation:
         > <The changes that you made>
         code:
         > <code line 1>
-        > <code line 2...>
-        
-        You should return one entry like above for each file you want to output, separated by '==='
-        """)
+        > <code line 2...>""")
     else:
         full_prompt += dedent("""
         
@@ -64,14 +64,14 @@ def send_iffy_edit(prompt: str, code: Dict[str, str], clipboard: bool = False) -
         > <the code to be output line 2>
         > <the code to be output, line n...>""")
 
-    max_tokens = confirm_send(full_prompt, yes=True)
+    max_tokens = confirm_send(full_prompt, yes=yes)
 
     response = openai.Completion.create(
         engine="text-davinci-003",
         prompt=full_prompt,
         max_tokens=max_tokens,
         n=1,
-        temperature=0.6,
+        temperature=0.5,
     )
 
     parsed = parse_resp(response)

codegpt/main.py

@@ -2,7 +2,6 @@
 import typer
 
 from codegpt import prompts
-from codegpt.prompts import PromptKeys
 
 from codegpt import gpt_interface as gpt
 from codegpt import files
@@ -18,15 +17,16 @@
 
 @app.command("do")
 def edit_file(
-    filenames: List[Path] = typer.Argument(
-        None, help="List of filenames to edit. If not provided, will prompt for input.",
+    instruction: str = typer.Argument(
+        ..., help="Instruction to edit the file(s). Keep it short! Wrap with quotes.",
     ),
-    instruction: str = typer.Option(
-        None, '--instruction', '-i', help="Instruction to edit the file(s). Keep it short!",
+    filenames: List[Path] = typer.Argument(
+        ..., help="List of filenames to edit. If not provided, will prompt for input.",
     ),
     backup: bool = typer.Option(
         False, "--backup", "-b", help="Whether to create a backup of the original file(s).",
     ),
+    yes: bool = typer.Option(False, "--yes", "-y", help="Don't ask for confirmation.",),
     ):
     """
     Edit one or more files using codegpt.
@@ -37,22 +37,24 @@ def edit_file(
     if not filenames:
         filenames = typer.prompt("Enter the filenames to edit, separated by spaces").split()
     code = files.load_text(filenames)
-    result = gpt.send_iffy_edit(instruction, code)
+    result = gpt.send_iffy_edit(instruction, code, yes=yes)
     files.write_text(result, backup)
     typer.secho("Done!", color=typer.colors.BRIGHT_BLUE)
 
 
 @app.command("quick")
 def quick_edit_file(
-    filenames: List[str],
-    option: PromptKeys = typer.Option(PromptKeys.COMMENT, "--option", "-o", case_sensitive=False),
-    backup=False,
+    option: str = typer.Argument(..., help=f"{{{'|'.join(prompts.prompts.keys())}}}"),
+    filenames: List[str] = typer.Argument(..., help="Enter the filenames to edit, separated by spaces"),
+    backup: bool = typer.Option(False, '--backup', '-b', help="Whether to create a backup of the original file(s)."),
     ):
     """
     Edit a file using codegpt's built in prompts
     """
+    if option not in prompts.prompts:
+        raise typer.BadParameter(f"{option} is not a valid option. Must be one of {list(prompts.prompts.keys())}")
     code = files.load_text(filenames)
-    result = gpt.send_iffy_edit(prompts.prompts[option.value], code)
+    result = gpt.send_iffy_edit(prompts.prompts[option], code)
     files.write_text(result, backup)
     typer.secho("done", color=typer.colors.BRIGHT_BLUE)
 

codegpt/prompts.py

@@ -1,14 +1,5 @@
-from enum import Enum
-
-class PromptKeys(str, Enum):
-    COMMENT = 'comment'
-    VARNAMES = 'varnames'
-    UGH = 'ugh'
-    BUGS = 'bugs'
-    VULNS = 'vulns'
-
 prompts = {
-    "comment": "Add or update comments according to the language's standards. Add or update function, module, or class level comments if they're appropriate.",
+    "comment": "Add or update comments according to the given language's standards. Add or update function, module, or class level comments if they're appropriate.",
     "varnames": "Change variable names, but nothing else, to make the code more readable. For example, instead of using 'x' and 'y', use 'width' and 'height'.",
     "ugh": "Do anything you can to make this code more readable. Add comments, change variable and function names, add whitespace, whatever. Add a readme to explain what the code does and where it could be improved.",
     "bugs": """Find any bugs you can, note them in comments prefixed with BUG:

codegpt/readme.md

@@ -16,7 +16,7 @@ pip install codegpt
 To use the CLI, simply run the following command:
 
 ```
-codegpt <command> <options
+codegpt <command> <instruction|choice> <filenames>
 ```
 
 ## Contributing

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "codegpt"
-version = "0.2.5"
+version = "0.2.7"
 license = "GPL-3.0-or-later"
 description = "A CLI tool for refactoring code using OpenAI's models"
 authors = ["John Partee"]