etc/tool/improve_code.py


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45


import sys, re
from pathlib import Path
from os import path

sys.path.append(str(Path(__file__).parent.parent.parent))

import g4f

def read_code(text):
    if match := re.search(r"```(python|py|)\n(?P<code>[\S\s]+?)\n```", text):
        return match.group("code")
    
path = input("Path: ")

with open(path, "r") as file:
    code = file.read()

prompt = f"""
Improve the code in this file:
```py
{code}
```
Don't remove anything.
Add typehints if possible.
Don't add any typehints to kwargs.
Don't remove license comments.
"""

print("Create code...")
response = []
for chunk in g4f.ChatCompletion.create(
    model=g4f.models.default,
    messages=[{"role": "user", "content": prompt}],
    timeout=300,
    stream=True
):
    response.append(chunk)
    print(chunk, end="", flush=True)
print()
response = "".join(response)

if code := read_code(response):
    with open(path, "w") as file:
        file.write(code)