holy shit add tests!!!

Author: morganpartee

codegpt/main.py

@@ -16,53 +16,99 @@
 
 app = typer.Typer()
 
+
 @app.command("do")
 def edit_file(
     instruction: str = typer.Argument(
-        ..., help="Instruction to edit the file(s). Keep it short! Wrap with quotes.",
+        ...,
+        help="Instruction to edit the file(s). Keep it short! Wrap with quotes.",
     ),
     filenames: List[Path] = typer.Argument(
-        [], help="List of filenames to edit. If not provided, will prompt for input.",
+        ["#CLIPBOARD"],
+        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).",
+        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.",
+    ),
+    raw_code: str = typer.Option(
+        None,
+        "--raw-code",
+        "-c",
+        help="Raw code to edit. Overrides filenames. Use quotes to wrap the code.",
+    ),
+    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.",
     ),
-    yes: bool = typer.Option(False, "--yes", "-y", help="Don't ask for confirmation.",),
-    raw_code: str = typer.Option(None, "--raw-code", "-c", help="Raw code to edit. Overrides filenames"),
-    json_out: bool = typer.Option(False, "--json-out", "-j", help="Output to raw json."),
-    ):
+):
     """
     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.
-    
-    Your code better be in git before you use this.
+
+    Your code better be in git before you use this. If the instruction is one of the quick prompt options (like 'comment' or 'docs'), it will do that prompt automatically. For more info, run 'codegpt quick --help'.
 
     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 not filenames and not raw_code:
-        raise typer.BadParameter("Either FILENAMES or --raw-code (-c) must be provided.")
+        raise typer.BadParameter(
+            "Either filenames or --raw-code (-c) must be provided."
+        )
+
     code = {"code": raw_code} if raw_code else files.load_text(filenames)
+
+    if instruction in prompts.prompts:
+        instruction = prompts.prompts[instruction]
+
     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)
 
+    if raw_out:
+        return 'result["code"]'
+
     files.write_text(result, backup)
     typer.secho("Done!", color=typer.colors.BRIGHT_BLUE)
 
 
 @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", "-b", help="Whether to create a backup of the original file(s).",
+    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).",
+    ),
+    yes: bool = typer.Option(
+        False,
+        "--yes",
+        "-y",
+        help="Don't ask for confirmation.",
+    ),
+    raw_code: str = typer.Option(
+        None, "--raw-code", "-c", help="Raw code to edit. Overrides filenames"
     ),
-    yes: bool = typer.Option(False, "--yes", "-y", help="Don't ask for confirmation.",),
-    raw_code: str = typer.Option(None, "--raw-code", "-c", help="Raw code to edit. Overrides filenames"),
     json_out: bool = typer.Option(False, "--json", "-j", help="Output in JSON format"),
-    ):
+):
     """
     Edit a file using codegpt's built in prompts.
 
@@ -75,13 +121,19 @@ def quick_edit_file(
     - vulns - Comment in code where the vulns are if GPT sees them (iffy)
     """
     if option not in prompts.prompts:
-        raise typer.BadParameter(f"{option} is not a valid option. Must be one of {list(prompts.prompts.keys())}")
+        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.")
+        raise typer.BadParameter(
+            "Either FILENAMES or --raw-code (-c) must be provided."
+        )
 
     code = {"code": raw_code} if raw_code else files.load_text(filenames)
-    result = gpt.send_iffy_edit(prompts.prompts[option], code, yes=yes, clipboard=bool(raw_code))
+    result = gpt.send_iffy_edit(
+        prompts.prompts[option], code, yes=yes, clipboard=bool(raw_code)
+    )
 
     if json_out:
         return json.dumps(result, sort_keys=True, indent=4)
@@ -97,7 +149,10 @@ def config():
     """
     # check if the secret key is already set in the environment variables
     if "OPENAI_SECRET_KEY" in os.environ:
-        typer.secho("OPENAI_SECRET_KEY is already set in the environment! You probably don't need this.", typer.colors.BRIGHT_BLUE)
+        typer.secho(
+            "OPENAI_SECRET_KEY is already set in the environment! You probably don't need this.",
+            typer.colors.BRIGHT_BLUE,
+        )
     else:
         typer.confirm(
             """

pyproject.toml

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

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+python_files = test_*.py
+testpaths = tests

tests/conftest.py

@@ -0,0 +1,24 @@
+import subprocess
+import sys
+from tqdm import tqdm
+
+
+import subprocess
+import sys
+from tqdm import tqdm
+
+
+def pytest_configure(config):
+    print("Building project...")
+    with tqdm(total=1, desc="Building project", file=sys.stdout) as pbar:
+        subprocess.run(["poetry", "build"])
+        pbar.update(1)
+    print("Installing project...")
+    with tqdm(total=1, desc="Installing project", file=sys.stdout) as pbar:
+        subprocess.run(["pip", "install", "--force", "--find-links=dist", "codegpt"])
+        pbar.update(1)
+    print("Installation complete.")
+
+
+if __name__ == "__main__":
+    pytest_configure({})

tests/test_edit_file.py

@@ -0,0 +1,46 @@
+import re
+from codegpt.main import edit_file
+
+
+def test_edit_file_add_comment():
+    # Set up the test inputs
+    instruction = "comment"
+    raw_code = "def foo():\n    print('Hello, world!')"
+
+    # Call the function being tested
+    result = edit_file(instruction, raw_code=raw_code, yes=True, raw_out=True)
+
+    # Use a regex to match the output of the function
+    assert "#" in result
+    assert "def foo():" in result
+    assert "print('Hello, world!')" in result
+
+
+def test_edit_file_reverse_string():
+    # Set up the test inputs and expected outputs
+    instruction = "reverse this string"
+    raw_code = "abcdef"
+    expected_output = {"code": "fedcba"}
+
+    # Call the function being tested
+    result = edit_file(instruction, raw_code=raw_code, yes=True, raw_out=True)
+
+    print(result)
+
+    # Verify that the function is behaving as expected
+    assert len(result) >= len(expected_output["code"])
+
+
+def test_edit_file_convert_to_uppercase():
+    # Set up the test inputs and expected outputs
+    instruction = "make this uppercase"
+    raw_code = "abcdef"
+    expected_output = "ABCDEF"
+
+    # Call the function being tested
+    result = edit_file(instruction, raw_code=raw_code, yes=True, raw_out=True)
+
+    print(result)
+
+    # Verify that the function is behaving as expected
+    assert len(result) >= len(expected_output)