mirror of
https://github.com/arc53/DocsGPT
synced 2024-11-03 23:15:37 +00:00
96 lines
2.8 KiB
Python
96 lines
2.8 KiB
Python
import ast
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import dotenv
|
|
from langchain.llms import OpenAI
|
|
from langchain.prompts import PromptTemplate
|
|
|
|
dotenv.load_dotenv()
|
|
|
|
ps = list(Path("inputs").glob("**/*.py"))
|
|
data = []
|
|
sources = []
|
|
for p in ps:
|
|
with open(p) as f:
|
|
data.append(f.read())
|
|
sources.append(p)
|
|
|
|
|
|
def get_functions_in_class(node):
|
|
functions = []
|
|
functions_code = []
|
|
for child in node.body:
|
|
if isinstance(child, ast.FunctionDef):
|
|
functions.append(child.name)
|
|
functions_code.append(ast.unparse(child))
|
|
|
|
return functions, functions_code
|
|
|
|
|
|
def get_classes_and_functions(source_code):
|
|
tree = ast.parse(source_code)
|
|
classes = {}
|
|
for node in tree.body:
|
|
if isinstance(node, ast.ClassDef):
|
|
class_name = node.name
|
|
function_name, function = get_functions_in_class(node)
|
|
# join function name and function code
|
|
functions = dict(zip(function_name, function))
|
|
classes[class_name] = functions
|
|
return classes
|
|
|
|
|
|
structure_dict = {}
|
|
c1 = 0
|
|
for code in data:
|
|
classes = get_classes_and_functions(ast.parse(code))
|
|
source = str(sources[c1])
|
|
structure_dict[source] = classes
|
|
c1 += 1
|
|
|
|
# save the structure dict as json
|
|
with open('structure_dict.json', 'w') as f:
|
|
json.dump(structure_dict, f)
|
|
|
|
if not Path("outputs").exists():
|
|
Path("outputs").mkdir()
|
|
|
|
c1 = len(structure_dict)
|
|
c2 = 0
|
|
for source, classes in structure_dict.items():
|
|
c2 += 1
|
|
print(f"Processing file {c2}/{c1}")
|
|
f1 = len(classes)
|
|
f2 = 0
|
|
for class_name, functions in classes.items():
|
|
f2 += 1
|
|
print(f"Processing class {f2}/{f1}")
|
|
source_w = source.replace("inputs/", "")
|
|
source_w = source_w.replace(".py", ".txt")
|
|
if not Path(f"outputs/{source_w}").exists():
|
|
with open(f"outputs/{source_w}", "w") as f:
|
|
f.write(f"Class: {class_name}")
|
|
else:
|
|
with open(f"outputs/{source_w}", "a") as f:
|
|
f.write(f"\n\nClass: {class_name}")
|
|
# append class name to the front
|
|
for function in functions:
|
|
b1 = len(functions)
|
|
b2 = 0
|
|
print(f"Processing function {b2}/{b1}")
|
|
b2 += 1
|
|
prompt = PromptTemplate(
|
|
input_variables=["code"],
|
|
template="Code: \n{code}, \nDocumentation: ",
|
|
)
|
|
llm = OpenAI(temperature=0)
|
|
response = llm(prompt.format(code=functions[function]))
|
|
|
|
if not Path(f"outputs/{source_w}").exists():
|
|
with open(f"outputs/{source_w}", "w") as f:
|
|
f.write(f"Function: {functions[function]}, \nDocumentation: {response}")
|
|
else:
|
|
with open(f"outputs/{source_w}", "a") as f:
|
|
f.write(f"\n\nFunction: {functions[function]}, \nDocumentation: {response}")
|