update and fix a lot of stuff, including tests

Author: morganpartee

.vscode/settings.json

@@ -2,9 +2,6 @@
     "cSpell.words": [
         "codegpt"
     ],
-    "python.testing.pytestArgs": [
-        "next"
-    ],
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true
 }

README.md

@@ -1,6 +1,6 @@
 # Codegpt
 
-## 0.2.12
+## 0.2.13
 
 A tool for using GPT just a little quicker. A nearly truly automated footgun. Learn how to revert with git before trying please.
 
@@ -27,7 +27,7 @@ Usage
 To try Codegpt, you can run the following command:
 
 ```bash
-codegpt do <instructions (quoted)> <filenames>
+codegpt do <instructions (quoted)> -f readme.md 
 ```
 
 It can do basically anything. Try handing in some files for context and telling it to generate something new - SQL queries, new features, documentation, whatever.

pyproject.toml

@@ -21,4 +21,24 @@ requires = ["poetry>=0.12"]
 build-backend = "poetry.masonry.api"
 
 [tool.poetry.scripts]
-codegpt = "codegpt.main:app"
+codegpt = "src.main:app"
+
+[tool.pytest]
+addopts = "--capture=fd --cov=my_module"
+
+# Directory where cache files should be stored
+cache_dir = ".pytest_cache"
+
+# List of plugins to load
+plugins = ["pytest-cov"]
+
+# List of files or directories to ignore during testing
+norecursedirs = ["venv"]
+
+# List of markers to exclude from test discovery
+markers = ["not_ready: mark tests that are not yet ready to be run"]
+
+[tool.pytest.ini_options]
+pythonpath = ["src"]
+testpaths = ["tests"]
+

pytest.ini

@@ -3,21 +3,23 @@
 import typer
 from typing import Dict
 from textwrap import dedent
-from codegpt.parse import parse_resp
+from parse import parse_resp
 
 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):
+def confirm_send(prompt, max_tokens=4000, yes=False, silent=False):
     tokens = nltk.word_tokenize(prompt)
 
     #! Yeah this math is BS, closeish though...
     max_tokens = round(max_tokens - (7 / 4) * len(tokens))
 
-    if yes:
+    if silent:
+        pass
+    elif yes:
         typer.secho(f"This prompt is {len(tokens)}ish tokens, GPT-3 can return {max_tokens}ish.", color=typer.colors.GREEN)
     else:
         typer.confirm(
@@ -65,7 +67,7 @@ def send_iffy_edit(prompt: str, code: Dict[str, str], clipboard: bool = False, y
         > <code line 1>
         > <code line n...>""")
 
-    max_tokens = confirm_send(full_prompt, yes=yes)
+    max_tokens = confirm_send(full_prompt, yes=yes, silent=clipboard)
 
     response = openai.Completion.create(
         engine="text-davinci-003",
@@ -75,8 +77,11 @@ def send_iffy_edit(prompt: str, code: Dict[str, str], clipboard: bool = False, y
         temperature=0.5,
     )
 
-    parsed = parse_resp(response)
-    
+    try:
+        parsed = parse_resp(response)
+    except KeyError as e:
+        print("ERROR: Response was malformed. Might still be usable, but is likely missing an explanation. Printing >\n")
+        print(response["choices"][0]["text"])
     return parsed[0] if clipboard else parsed
 
 

codegpt/__init__.py renamed to src/__init__.py

@@ -1,13 +1,15 @@
 import os
 import typer
 import json
+import logging
 
-from codegpt import prompts
 
-from codegpt import gpt_interface as gpt
-from codegpt import files
+import gpt_interface as gpt
 
-from typing import List
+import prompts
+import files
+
+from typing import List, Optional
 from pathlib import Path
 
 app = typer.Typer(
@@ -23,10 +25,6 @@ def edit_file(
         ...,
         help="Instruction to edit the file(s). Keep it short! Wrap with quotes.",
     ),
-    filenames: List[Path] = typer.Argument(
-        ["#CLIPBOARD"],
-        help="List of filenames to edit. If not provided, will prompt for input.",
-    ),
     backup: bool = typer.Option(
         False,
         "--backup",
@@ -54,7 +52,8 @@ def edit_file(
         "-r",
         help="Output the raw 'code' from the response and exit the function.",
     ),
-):
+        filenames: Optional[List[Path]] = typer.Argument(None, help="File(s) to edit or for context."),
+    ):
     """
     Do something given some code for context. Asking for documents, queries, etc. should work okay. Edits are iffy, but work a lot of the time.
 
@@ -63,10 +62,13 @@ def edit_file(
     FILENAMES: list of filenames to edit. If not provided, will prompt for input.
     INSTRUCTION: the instruction to edit the file(s). Keep it short!
     """
+    
+    if raw_out or json_out:
+        logging.basicConfig(level=logging.CRITICAL)
 
     if not filenames and not raw_code:
         raise typer.BadParameter(
-            "Either filenames or --raw-code (-c) must be provided."
+            "Either --filenames (-f) or --raw-code (-c) must be provided."
         )
 
     code = {"code": raw_code} if raw_code else files.load_text(filenames)
@@ -77,10 +79,12 @@ def edit_file(
     result = gpt.send_iffy_edit(instruction, code, yes=yes, clipboard=bool(raw_code))
 
     if json_out:
-        return json.dumps(result, sort_keys=True, indent=4)
+        print(json.dumps(result, sort_keys=True, indent=4))
+        return
 
     if raw_out:
-        return 'result["code"]'
+        print(result['code'])
+        return
 
     files.write_text(result, backup)
     typer.secho("Done!", color=typer.colors.BRIGHT_BLUE)
@@ -89,9 +93,6 @@ def edit_file(
 @app.command("quick")
 def quick_edit_file(
     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",
@@ -105,9 +106,21 @@ def quick_edit_file(
         help="Don't ask for confirmation.",
     ),
     raw_code: str = typer.Option(
-        None, "--raw-code", "-c", help="Raw code to edit. Overrides filenames"
+        None,
+        "--raw-code",
+        "-c",
+        help="Raw code to edit. Overrides filenames. Use quotes to wrap the code.",
     ),
-    json_out: bool = typer.Option(False, "--json", "-j", help="Output in JSON format"),
+    json_out: bool = typer.Option(
+        False, "--json-out", "-j", help="Output the response in raw json format."
+    ),
+    raw_out: bool = typer.Option(
+        False,
+        "--raw-out",
+        "-r",
+        help="Output the raw 'code' from the response and exit the function.",
+    ),
+    filenames: Optional[List[Path]] = typer.Argument(None, help="File(s) to edit or for context."),
 ):
     """
     Edit a file using codegpt's built in prompts.
@@ -124,7 +137,7 @@ def quick_edit_file(
         raise typer.BadParameter(
             f"{option} is not a valid option. Must be one of {list(prompts.prompts.keys())}"
         )
-
+    
     if not filenames and not raw_code:
         raise typer.BadParameter(
             "Either FILENAMES or --raw-code (-c) must be provided."
@@ -136,7 +149,12 @@ def quick_edit_file(
     )
 
     if json_out:
-        return json.dumps(result, sort_keys=True, indent=4)
+        print(json.dumps(result, sort_keys=True, indent=4))
+        return
+
+    if raw_out: 
+        print(result['code'])
+        return
 
     files.write_text(result, backup)
     typer.secho("Done!", color=typer.colors.BRIGHT_BLUE)

codegpt/__main__.py renamed to src/__main__.py

@@ -1,5 +1,5 @@
 def parse_resp(response:dict):
-    resp = response["choices"][0]["text"].splitlines()
+    resp = response["choices"][0]["text"].strip().splitlines()
 
     # Initialize an empty list to hold the dictionaries
     out = []
@@ -35,4 +35,8 @@ def parse_resp(response:dict):
     # Add the final dictionary to the output list
     out.append(curr_dict)
 
+    # Backtop just in case this fails. Tests don't tend to use whole code, so it gets weird.
+    if 'code' not in out:
+        out = [{'code': response["choices"][0]["text"]}]
+
     return out