Skip to content

Commit f9183d3

Browse files
authored
Merge pull request #11 from morganpartee/main
catchup
2 parents 904f6e2 + 672a8f7 commit f9183d3

File tree

8 files changed

+212
-146
lines changed

8 files changed

+212
-146
lines changed

.vscode/settings.json

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,6 @@
22
"cSpell.words": [
33
"codegpt"
44
],
5-
"python.testing.pytestArgs": [
6-
"next"
7-
],
85
"python.testing.unittestEnabled": false,
96
"python.testing.pytestEnabled": true
107
}

README.md

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
# Codegpt
22

3-
## 0.2.10
3+
## 0.2.15
44

55
A tool for using GPT just a little quicker. A nearly truly automated footgun. Learn how to revert with git before trying please.
66

@@ -27,10 +27,10 @@ Usage
2727
To try Codegpt, you can run the following command:
2828

2929
```bash
30-
codegpt do <instructions (quoted)> <filenames>
30+
codegpt do <instructions (quoted)> -f readme.md
3131
```
3232

33-
It will prompt you for directions to follow, and it'll do whatever you want. Write new docs, add comments to code, just make sure to ask it to "edit" or "comment" or whatever so it knows to change files or not.
33+
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.
3434

3535
Or use the quick command to do some neat stuff, like:
3636

codegpt/gpt_interface.py

Lines changed: 24 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,15 @@
1111
typer.secho("Downloading punkt for nltk... Only once!", fg=typer.colors.GREEN, bold=True)
1212
nltk.download('punkt', quiet=True)
1313

14-
def confirm_send(prompt, max_tokens=4000, yes=False):
14+
def confirm_send(prompt, max_tokens=4000, yes=False, silent=False):
1515
tokens = nltk.word_tokenize(prompt)
1616

1717
#! Yeah this math is BS, closeish though...
1818
max_tokens = round(max_tokens - (7 / 4) * len(tokens))
1919

20-
if yes:
20+
if silent:
21+
pass
22+
elif yes:
2123
typer.secho(f"This prompt is {len(tokens)}ish tokens, GPT-3 can return {max_tokens}ish.", color=typer.colors.GREEN)
2224
else:
2325
typer.confirm(
@@ -42,30 +44,38 @@ def send_iffy_edit(prompt: str, code: Dict[str, str], clipboard: bool = False, y
4244
if clipboard:
4345
full_prompt += dedent("""
4446
45-
Answer in the following format:
47+
Answer in the following format, using '> ' to indent values:
4648
4749
explanation:
4850
> <The changes that you made>
4951
code:
5052
> <the code to be output line 1>
51-
> <the code to be output, line n...>""")
53+
> <the code to be output, line n...>
54+
55+
You must include an explanation of what you did, and the code to be output, regardless of the format or file.
56+
57+
OUTPUT:""")
5258

5359
else:
5460
full_prompt += dedent("""
5561
56-
You may only output complete files.
57-
58-
If you add or modify a file, return it in this exact format:
62+
You may only send me complete files.
5963
64+
For each file, return one block in this format, using '> ' to indent values:
65+
6066
filename:
6167
> <the filename to be output>
6268
explanation:
6369
> <The changes that you made>
6470
code:
6571
> <code line 1>
66-
> <code line n...>""")
72+
> <code line n...>
73+
74+
You must include the filename, an explanation of what you did, and the code for the file to be output, regardless of the format or file.
75+
76+
OUTPUT:""")
6777

68-
max_tokens = confirm_send(full_prompt, yes=yes)
78+
max_tokens = confirm_send(full_prompt, yes=yes, silent=clipboard)
6979

7080
response = openai.Completion.create(
7181
engine="text-davinci-003",
@@ -75,8 +85,11 @@ def send_iffy_edit(prompt: str, code: Dict[str, str], clipboard: bool = False, y
7585
temperature=0.5,
7686
)
7787

78-
parsed = parse_resp(response)
79-
88+
try:
89+
parsed = parse_resp(response)
90+
except KeyError as e:
91+
print("ERROR: Response was malformed. Might still be usable, but is likely missing an explanation. Printing >\n")
92+
print(response["choices"][0]["text"])
8093
return parsed[0] if clipboard else parsed
8194

8295

codegpt/main.py

Lines changed: 100 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
import os
22
import typer
33
import json
4+
import logging
45

5-
from codegpt import prompts
66

77
from codegpt import gpt_interface as gpt
8+
9+
from codegpt import prompts
810
from codegpt import files
911

10-
from typing import List
12+
from typing import List, Optional
1113
from pathlib import Path
1214

1315
app = typer.Typer(
@@ -16,37 +18,73 @@
1618

1719
app = typer.Typer()
1820

21+
1922
@app.command("do")
2023
def edit_file(
2124
instruction: str = typer.Argument(
22-
..., help="Instruction to edit the file(s). Keep it short! Wrap with quotes.",
23-
),
24-
filenames: List[Path] = typer.Argument(
25-
[], help="List of filenames to edit. If not provided, will prompt for input.",
25+
...,
26+
help="Instruction to edit the file(s). Keep it short! Wrap with quotes.",
2627
),
2728
backup: bool = typer.Option(
28-
False, "--backup", "-b", help="Whether to create a backup of the original file(s).",
29+
False,
30+
"--backup",
31+
"-b",
32+
help="Whether to create a backup of the original file(s).",
33+
),
34+
yes: bool = typer.Option(
35+
False,
36+
"--yes",
37+
"-y",
38+
help="Don't ask for confirmation.",
39+
),
40+
raw_code: str = typer.Option(
41+
None,
42+
"--raw-code",
43+
"-c",
44+
help="Raw code to edit. Overrides filenames. Use quotes to wrap the code.",
45+
),
46+
json_out: bool = typer.Option(
47+
False, "--json-out", "-j", help="Output the response in raw json format."
2948
),
30-
yes: bool = typer.Option(False, "--yes", "-y", help="Don't ask for confirmation.",),
31-
raw_code: str = typer.Option(None, "--raw-code", "-c", help="Raw code to edit. Overrides filenames"),
32-
json_out: bool = typer.Option(False, "--json-out", "-j", help="Output to raw json."),
49+
raw_out: bool = typer.Option(
50+
False,
51+
"--raw-out",
52+
"-r",
53+
help="Output the raw 'code' from the response and exit the function.",
54+
),
55+
filenames: Optional[List[Path]] = typer.Argument(None, help="File(s) to edit or for context."),
3356
):
3457
"""
3558
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.
36-
37-
Your code better be in git before you use this.
59+
60+
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'.
3861
3962
FILENAMES: list of filenames to edit. If not provided, will prompt for input.
4063
INSTRUCTION: the instruction to edit the file(s). Keep it short!
4164
"""
65+
66+
if raw_out or json_out:
67+
logging.basicConfig(level=logging.CRITICAL)
4268

4369
if not filenames and not raw_code:
44-
raise typer.BadParameter("Either FILENAMES or --raw-code (-c) must be provided.")
70+
raise typer.BadParameter(
71+
"Either --filenames (-f) or --raw-code (-c) must be provided."
72+
)
73+
4574
code = {"code": raw_code} if raw_code else files.load_text(filenames)
75+
76+
if instruction in prompts.prompts:
77+
instruction = prompts.prompts[instruction]
78+
4679
result = gpt.send_iffy_edit(instruction, code, yes=yes, clipboard=bool(raw_code))
4780

4881
if json_out:
49-
return json.dumps(result, sort_keys=True, indent=4)
82+
print(json.dumps(result, sort_keys=True, indent=4))
83+
return
84+
85+
if raw_out:
86+
print(result['code'])
87+
return
5088

5189
files.write_text(result, backup)
5290
typer.secho("Done!", color=typer.colors.BRIGHT_BLUE)
@@ -55,14 +93,35 @@ def edit_file(
5593
@app.command("quick")
5694
def quick_edit_file(
5795
option: str = typer.Argument(..., help=f"{{{'|'.join(prompts.prompts.keys())}}}"),
58-
filenames: List[str] = typer.Argument(..., help="Enter the filenames to edit, separated by spaces"),
59-
backup: bool = typer.Option(
60-
False, "--backup", "-b", help="Whether to create a backup of the original file(s).",
96+
backup: bool = typer.Option(
97+
False,
98+
"--backup",
99+
"-b",
100+
help="Whether to create a backup of the original file(s).",
61101
),
62-
yes: bool = typer.Option(False, "--yes", "-y", help="Don't ask for confirmation.",),
63-
raw_code: str = typer.Option(None, "--raw-code", "-c", help="Raw code to edit. Overrides filenames"),
64-
json_out: bool = typer.Option(False, "--json", "-j", help="Output in JSON format"),
65-
):
102+
yes: bool = typer.Option(
103+
False,
104+
"--yes",
105+
"-y",
106+
help="Don't ask for confirmation.",
107+
),
108+
raw_code: str = typer.Option(
109+
None,
110+
"--raw-code",
111+
"-c",
112+
help="Raw code to edit. Overrides filenames. Use quotes to wrap the code.",
113+
),
114+
json_out: bool = typer.Option(
115+
False, "--json-out", "-j", help="Output the response in raw json format."
116+
),
117+
raw_out: bool = typer.Option(
118+
False,
119+
"--raw-out",
120+
"-r",
121+
help="Output the raw 'code' from the response and exit the function.",
122+
),
123+
filenames: Optional[List[Path]] = typer.Argument(None, help="File(s) to edit or for context."),
124+
):
66125
"""
67126
Edit a file using codegpt's built in prompts.
68127
@@ -75,16 +134,27 @@ def quick_edit_file(
75134
- vulns - Comment in code where the vulns are if GPT sees them (iffy)
76135
"""
77136
if option not in prompts.prompts:
78-
raise typer.BadParameter(f"{option} is not a valid option. Must be one of {list(prompts.prompts.keys())}")
79-
137+
raise typer.BadParameter(
138+
f"{option} is not a valid option. Must be one of {list(prompts.prompts.keys())}"
139+
)
140+
80141
if not filenames and not raw_code:
81-
raise typer.BadParameter("Either FILENAMES or --raw-code (-c) must be provided.")
142+
raise typer.BadParameter(
143+
"Either FILENAMES or --raw-code (-c) must be provided."
144+
)
82145

83146
code = {"code": raw_code} if raw_code else files.load_text(filenames)
84-
result = gpt.send_iffy_edit(prompts.prompts[option], code, yes=yes, clipboard=bool(raw_code))
147+
result = gpt.send_iffy_edit(
148+
prompts.prompts[option], code, yes=yes, clipboard=bool(raw_code)
149+
)
85150

86151
if json_out:
87-
return json.dumps(result, sort_keys=True, indent=4)
152+
print(json.dumps(result, sort_keys=True, indent=4))
153+
return
154+
155+
if raw_out:
156+
print(result['code'])
157+
return
88158

89159
files.write_text(result, backup)
90160
typer.secho("Done!", color=typer.colors.BRIGHT_BLUE)
@@ -97,7 +167,10 @@ def config():
97167
"""
98168
# check if the secret key is already set in the environment variables
99169
if "OPENAI_SECRET_KEY" in os.environ:
100-
typer.secho("OPENAI_SECRET_KEY is already set in the environment! You probably don't need this.", typer.colors.BRIGHT_BLUE)
170+
typer.secho(
171+
"OPENAI_SECRET_KEY is already set in the environment! You probably don't need this.",
172+
typer.colors.BRIGHT_BLUE,
173+
)
101174
else:
102175
typer.confirm(
103176
"""

codegpt/parse.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
def parse_resp(response:dict):
2-
resp = response["choices"][0]["text"].splitlines()
2+
resp = response["choices"][0]["text"].strip().splitlines()
33

44
# Initialize an empty list to hold the dictionaries
55
out = []
@@ -35,4 +35,8 @@ def parse_resp(response:dict):
3535
# Add the final dictionary to the output list
3636
out.append(curr_dict)
3737

38+
# Backtop just in case this fails. Tests don't tend to use whole code, so it gets weird.
39+
if 'code' not in out:
40+
out = [{'code': response["choices"][0]["text"]}]
41+
3842
return out

pyproject.toml

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[tool.poetry]
22
name = "codegpt"
3-
version = "0.2.10"
3+
version = "0.2.15"
44
license = "GPL-3.0-or-later"
55
description = "A CLI tool for refactoring code using OpenAI's models"
66
authors = ["John Partee"]
@@ -22,3 +22,23 @@ build-backend = "poetry.masonry.api"
2222

2323
[tool.poetry.scripts]
2424
codegpt = "codegpt.main:app"
25+
26+
[tool.pytest]
27+
addopts = "--capture=fd --cov=my_module"
28+
29+
# Directory where cache files should be stored
30+
cache_dir = ".pytest_cache"
31+
32+
# List of plugins to load
33+
plugins = ["pytest-cov"]
34+
35+
# List of files or directories to ignore during testing
36+
norecursedirs = ["venv"]
37+
38+
# List of markers to exclude from test discovery
39+
markers = ["not_ready: mark tests that are not yet ready to be run"]
40+
41+
[tool.pytest.ini_options]
42+
pythonpath = ["codegpt"]
43+
testpaths = ["tests"]
44+

tests/test_edit_file.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
from main import app
2+
3+
from typer.testing import CliRunner
4+
5+
run = CliRunner()
6+
7+
8+
def test_edit_file_add_comment():
9+
# Set up the test inputs
10+
instruction = "Add one comment with a '#'"
11+
raw_code = "def foo():\n print('Hello, world!')"
12+
13+
# Call the function being tested
14+
result = run.invoke(app, ['do', instruction, "--yes", "--raw-out", "--raw-code", f"'{raw_code}'"])
15+
16+
# Use a regex to match the output of the function
17+
assert "#" in result.stdout
18+
assert "def foo():" in result.stdout
19+
assert "print('Hello, world!')" in result.stdout
20+
21+
22+
def test_varnames():
23+
# Set up the test inputs and expected outputs
24+
instruction = "varnames"
25+
raw_code = """
26+
v0 = 0.2
27+
a = 2
28+
n = 21 # No of t values for plotting
29+
30+
t = np.linspace(0, 2, n+1)
31+
s = v0*t + 0.5*a*t**2
32+
"""
33+
34+
# Call the function being tested
35+
result = run.invoke(app, ['do', instruction, "--yes", "--raw-out", "--raw-code", f"'{raw_code}'"])
36+
37+
assert 'v0 = 0.2' in result.stdout
38+
assert 'a = 2' not in result.stdout
39+
assert 'n = 21' not in result.stdout
40+
assert 't = np.linspace(0, 2, n+1)' not in result.stdout
41+
assert's = v0*t + 0.5*a*t**2' not in result.stdout
42+
43+
if __name__ == "__main__":
44+
import pytest
45+
46+
pytest.main([__file__])

0 commit comments

Comments
 (0)