python-docstring

#!/usr/bin/python
importast
importos
importrequests
importsys
importjson
importtempfile
importsubprocess
importhashlib
importpprint
importjson


classOpenAIProvider:
def__init__(self):
self.api_key=os.getenv("AI_API_KEY")
ifnotself.api_key:
raiseEnvironmentError("AI_API_KEY not set in environment.")
self.endpoint="https://api.openai.com/v1/chat/completions"
self.model=os.getenv("OPENAI_MODEL", "gpt-4o-mini")

defimprove_text(self, prompt: str, text: str) ->str:
headers= {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
 }

body= {
"model": self.model,
"messages": [
 {"role": "system", "content": prompt},
 {"role": "user", "content": text},
 ],
"temperature": 0.3,
 }

response=requests.post(self.endpoint, json=body, headers=headers, timeout=60)
ifresponse.status_code==200:
returnresponse.json()["choices"][0]["message"]["content"].strip()

raiseException(
f"OpenAI API call failed: {response.status_code} - {response.text}"
 )


defnormalize_docstring(docstring):
"""Normalize docstring to strip leading/trailing spaces, remove markdown, and preserve multi-line format."""
ifnotdocstring:
return""

# Extract everything between triple quotes
# This assumes the docstring is surrounded by triple quotes
start=docstring.find('"""') +3# Skip past the first triple quotes
end=docstring.rfind('"""') # Find the last triple quotes

# Extract the content between the triple quotes
ifstart>2andend>start:
docstring=docstring[start:end].strip() # Remove the triple quotes and extra spaces

returndocstring


defget_functions_and_classes_without_docstrings(tree):
"""Parse the file and return functions and classes without docstrings, 
 considering top-level functions and functions inside a class only."""

functions_and_classes_without_docstrings= []

# We need to track the current class level to identify nested functions
class_stack= []

fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
# Determine the level of indentation
indentation_level=node.col_offset

# Debugging: Print function name, indentation level, and docstring presence
docstring=ast.get_docstring(node)
docstring_status="No docstring"ifdocstringisNoneelse"Has docstring"

# Skip functions with indentation level of 8 (nested functions inside other functions)
ifindentation_level==8:
# print(f"Skipping nested function: {node.name}")
continue
else:
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('function', node))

elifisinstance(node, ast.ClassDef):
# Add the class to the class stack, as we're inside a class now
class_stack.append(node)

# Check if the class has a docstring
docstring=ast.get_docstring(node)
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('class', node))

# Pop the class stack when leaving a class body
ifisinstance(node, ast.ClassDef):
class_stack.pop()

returnfunctions_and_classes_without_docstrings


defgenerate_docstring(object_code, openai_provider, cache, object_name, object_type, filename):
"""Generate a docstring using OpenAIProvider."""
code_digest=hashlib.sha256(object_code.encode()).hexdigest()
entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()
ifentryincacheandcode_digest==cache[entry]['last_hash']:
print(f"Using cached docstring for {object_type}: {object_name}")
returncache[entry]['docstring'] # Return cached docstring if it exists

ifentryincacheandcode_digest!=cache[entry]['last_hash']:
print(f"Has for {object_type}: {object_name} changed. Getting new docstring.")

prompt=f"""
 You will be provided with a Python {object_type} (function or class),
 including its code. Based on this, your task is to generate a Python docstring for it.
 Use the function signature and body to infer the purpose of the function/class,
 the arguments it takes, the return value, and any exceptions it may raise.

 Follow these guidelines to create the docstring:

 1. Summary: Provide a concise summary of the function/class's purpose.
 Focus on what the function does and its main goal.
 2. Arguments: List the parameters, their types, and a brief description
 of what each one represents.
 3. Return: If the function has a return value, describe the return type
 and what it represents. If there's no return, OMIT THE SECTION.
 4. Exceptions: If the function raises any exceptions, list them with descriptions.
 If no exceptions are raised, OMIT THE SECTION.
 5. Side Effects (if applicable): If the function has side effects
 (e.g., modifies global state, interacts with external services), mention them.
 OMIT THE SECTION if it is not lear in the code.
 6. Algorithm or Key Logic (optional): If the function is complex,
 provide a high-level outline of the logic or algorithm involved.
 OMIT THE SECTION if it is not lear in the code.

 Ensure the docstring follows Python's standard docstring conventions and provides
 just enough detail to make the function understandable and usable without overwhelming the reader.
 """

# Get the docstring from OpenAI
docstring=openai_provider.improve_text(prompt, object_code)

# Normalize the docstring by stripping leading/trailing whitespace
normalized_docstring=normalize_docstring(docstring)

# If the docstring is empty, set a fallback docstring
ifnotnormalized_docstring:
print(f"Warning: Generated docstring for {object_type} '{object_name}' is empty. Using a fallback docstring.")
normalized_docstring=f"TODO: Add a docstring for the {object_type} '{object_name}'."

tree=ast.parse(object_code)
indentation_level=0
fornodeinast.walk(tree):
ifisinstance(node, (ast.ClassDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+4# Add 4 spaces for method body
break
ifisinstance(node, (ast.FunctionDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+8# Add 4 spaces for method body
break

# Split the normalized docstring into lines and prefix each with the correct indentation
docstring_lines=normalized_docstring.splitlines()
indented_docstring="\n".join([f"{' '*indentation_level}{line}"forlineindocstring_lines])

# Cache the docstring
cache[entry] = {}
cache[entry]['docstring'] =indented_docstring
cache[entry]['last_hash'] =hashlib.sha256(object_code.encode()).hexdigest()

returnindented_docstring

defvalidate_syntax(code):
"""Validate Python code syntax by writing to a temporary file and checking it."""
try:
# Create a temporary file to check the syntax
withtempfile.NamedTemporaryFile(delete=False, mode='w', suffix=".py") astemp_file:
temp_file.write(code)
temp_filename=temp_file.name

# Run python -m py_compile to check syntax
result=subprocess.run(
 ['python', '-m', 'py_compile', temp_filename],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
 )

ifresult.returncode!=0:
print(code)

os.remove(temp_filename) # Remove the temporary file after checking

ifresult.returncode!=0:
print(f"Syntax error: {result.stderr.decode()}")
returnFalse
returnTrue

exceptExceptionase:
print(f"Error validating syntax: {e}")
returnFalse


defprocess_directory(directory, openai_provider, cache, cache_file):
"""Process all Python files in a directory and its subdirectories."""
forroot, dirs, filesinos.walk(directory):
forfileinfiles:
iffile.endswith(".py"): # Process only Python files
filepath=os.path.join(root, file)
process_file(filepath, openai_provider, cache, cache_file)


defload_cache(cache_file):
"""Load the cache from a file."""
ifos.path.exists(cache_file):
withopen(cache_file, 'r') asf:
returnjson.load(f)
return {}


defsave_cache(cache, cache_file):
"""Save the cache to a file."""
withopen(cache_file, 'w') asf:
json.dump(cache, f, indent=4)

defget_next(tree):
"""Get the next function or class without a docstring."""
# Assuming this function returns the next function/class that lacks a docstring.
# Return None when all functions/classes are documented
objects_without_docstrings=get_functions_and_classes_without_docstrings(tree)
ifobjects_without_docstrings:
returnobjects_without_docstrings[0] # Return the next object
returnNone# No more objects to process

defget_object_node(tree, object_name):
"""
 Traverses the AST to find and return the node for the object (function or class)
 with the specified name.

 Args:
 tree (ast.Module): The AST module to search through.
 object_name (str): The name of the function or class to find.

 Returns:
 ast.FunctionDef | ast.ClassDef: The AST node for the function or class.
 """
fornodeinast.walk(tree):
# Check for function or class definitions by name
ifisinstance(node, (ast.FunctionDef, ast.ClassDef)) andnode.name==object_name:
returnnode
returnNone

defprocess_file(filename, openai_provider, cache, cache_file):
"""Process the file: find functions, classes, generate docstrings, and save the result."""
print(f"Processing file: {filename}")

# Reread the file to ensure we are working with the latest version of the code
withopen(filename, 'r') asfile:
code=file.read()

# Start an infinite loop that will keep processing until there are no more items
whileTrue:
tree=ast.parse(code)
# Get the next function or class that doesn't have a docstring
next_object=get_next(tree)
ifnotnext_object:
print(f"All functions and classes have docstrings in {filename}.")
break# Exit the loop if there are no more functions or classes to document

object_type, object_node=next_object
# Get the object code (e.g., function or class)
object_code=ast.unparse(object_node)
object_name=object_node.name
print(f"Generating docstring for {object_type}: {object_name} in {filename}")

# Check if the function or class already has a docstring
ifast.get_docstring(object_node):
print(f"Skipping {object_type}: {object_name}, it already has a docstring.")
continue

# Generate the docstring using OpenAI or another provider
docstring=generate_docstring(object_code, openai_provider, cache, object_name, object_type, filename)

ifdocstring:
code=insert_docstring_to_object(code, object_node, docstring, object_type)

# Re-read the object code after the docstring insertion
updated_tree=ast.parse(code)
updated_object_node=get_object_node(updated_tree, object_name)
updated_object_code=ast.unparse(updated_object_node)

entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()

# Save the SHA-1 hash of the updated object_code into the cache
cache[entry]['last_hash'] =hashlib.sha1(updated_object_code.encode()).hexdigest()

# Save cache after processing each object
save_cache(cache, cache_file)
print(f"Docstrings added successfully to {object_type}: {object_name} in {filename}.")
else:
print(f"Something wrong with the docstring generation.")

# Validate Python syntax before saving
ifvalidate_syntax(code):
# Write the updated code back to the file
withopen(filename, 'w') asfile: 
file.write(code) 
else:
print(f"Syntax error in {filename}, skipping save.")

deffix_indentation(code):
"""Fix the indentation of the function signature and the code after the docstring."""

# Split the code into lines
code_lines=code.splitlines()

# Find the index where the function signature ends (the first occurrence of ':')
def_index=next(ifori, lineinenumerate(code_lines) ifline.strip().startswith('def '))

# Add 4 spaces to the function signature
code_lines[def_index] =' '+code_lines[def_index].lstrip()

# Now, we need to add indentation for all lines after the docstring
in_docstring=False
foriinrange(def_index+1, len(code_lines)):
line=code_lines[i]

# Detect the start and end of the docstring
ifline.strip().startswith('"""'):
ifin_docstring:
# End of the docstring
in_docstring=False
else:
# Start of the docstring
in_docstring=True

# Indent the line after the docstring
ifin_docstringisFalse:
code_lines[i] =' '+line# Add 4 spaces to the body of the function

# Rejoin the lines into the final code
updated_code="\n".join(code_lines)

returnupdated_code


defremove_duplicate_functions(code):
"""Remove duplicate function definitions based on the function signature."""
# Parse the code into an AST
tree=None
try:
tree=ast.parse(code)
exceptExceptionase:
print(code)
returncode

# Track function signatures
function_signatures=set()
functions_to_remove= []

# Walk through the AST and find function definitions
fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
function_signature=ast.unparse(node)

# If we have seen this function signature before, mark it for removal
iffunction_signatureinfunction_signatures:
functions_to_remove.append(node)
else:
function_signatures.add(function_signature)

# Remove duplicate functions from the code
forfunc_nodeinfunctions_to_remove:
# Removing the function by filtering out the lines corresponding to that function
code=remove_function_from_code(code, func_node)

returncode


defremove_function_from_code(code, function_node):
"""Remove the function from the code."""
# Split the code into lines
code_lines=code.splitlines()

# Get the start and end line numbers of the function
start_line=function_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(ast.unparse(function_node).splitlines())

# Remove the lines of the function from the code
return"\n".join(code_lines[:start_line] +code_lines[end_line+1:])


defextract_main_block(code):
"""Extracts everything from 'if __name__ == "__main__":' to the end of the file."""
lines=code.splitlines() # Split the code into lines
main_start_index=-1# Initialize the index

# Find the line number where '__name__' starts
fori, lineinenumerate(lines):
ifline.strip().startswith('if __name__ == "__main__":'):
main_start_index=i
break

ifmain_start_index==-1:
# print("No '__name__ == \"__main__\"' found in the code.")
return""

# Extract all lines starting from the '__name__' line to the end
return"\n".join(lines[main_start_index:])


definsert_docstring_to_object(code, object_node, docstring, object_type):
main=extract_main_block(code)
"""Insert the generated docstring into the function/class in the correct place."""
object_code=ast.unparse(object_node)

# Get the indentation level of the function or class signature using col_offset
indentation_level=' '* (object_node.col_offset+4) # Add 4 spaces for the body

# Check if the function or class already has a docstring, if so, remove it
existing_docstring=ast.get_docstring(object_node)
ifexisting_docstring:
# Remove the existing docstring from the object_code
object_code=object_code.replace(f'"""{existing_docstring}"""', '')

# The docstring is already properly formatted and indented, so we use it directly
docstring_code=f'{indentation_level}"""\n{docstring}\n{indentation_level}"""'

# Find the first occurrence of ':' to determine where the signature ends
signature_end=object_code.find(':')

func_start_lines=object_code[:signature_end+1]
func_end_lines=object_code[signature_end+1:]

# Insert the docstring immediately after the signature (before the body of the function/class)
object_code_with_docstring=func_start_lines+f'\n{docstring_code}\n'+func_end_lines

ifobject_node.col_offset>0:
object_code_with_docstring=fix_indentation(object_code_with_docstring)

# Get the start and end line numbers of the object_node
start_line=object_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(object_code.splitlines()) -1

# Split the original code into lines
code_lines=code.splitlines()

print("================================")
print(pprint.pprint(code_lines[:start_line]))
print("================================")

# Replace the lines between start_line and end_line with the new code containing the docstring
updated_code_lines= (
code_lines[:start_line] +
 [object_code_with_docstring] +
code_lines[end_line+1:]
 )

# Join the lines back into a single string
updated_code="\n".join(updated_code_lines)
updated_code=remove_duplicate_functions(updated_code)
main_after_update=extract_main_block(updated_code)

ifmain!=main_after_update:
updated_code=updated_code+"\n\n"+main


# Ensure a line break after the inserted docstring for readability
ifnotupdated_code.endswith('\n'):
updated_code+='\n'

returnupdated_code


if__name__=='__main__':
# Get the directory or file from the command line argument
iflen(sys.argv) !=2:
print("Usage: python add_docstrings.py <python_file_or_directory>")
sys.exit(1)

path=sys.argv[1]

# Initialize the OpenAIProvider
openai_provider=OpenAIProvider()

# Load the cache from the cache file (if it exists)
cache_file="docstring_cache.json"
cache=load_cache(cache_file)

ifos.path.isdir(path): # If it's a directory, process it recursively
process_directory(path, openai_provider, cache, cache_file)
elifos.path.isfile(path): # If it's a single file, process it
process_file(path, openai_provider, cache, cache_file)
else:
print(f"The provided path '{path}' is not valid.")
sys.exit(1)

# Save the updated cache to the file
save_cache(cache, cache_file)
#!/usr/bin/python
importast
importos
importrequests
importsys
importjson
importtempfile
importsubprocess
importhashlib
importpprint
importjson


classOpenAIProvider:
def__init__(self):
self.api_key=os.getenv("AI_API_KEY")
ifnotself.api_key:
raiseEnvironmentError("AI_API_KEY not set in environment.")
self.endpoint="https://api.openai.com/v1/chat/completions"
self.model=os.getenv("OPENAI_MODEL", "gpt-4o-mini")

defimprove_text(self, prompt: str, text: str) ->str:
headers= {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
 }

body= {
"model": self.model,
"messages": [
 {"role": "system", "content": prompt},
 {"role": "user", "content": text},
 ],
"temperature": 0.3,
 }

response=requests.post(self.endpoint, json=body, headers=headers, timeout=60)
ifresponse.status_code==200:
returnresponse.json()["choices"][0]["message"]["content"].strip()

raiseException(
f"OpenAI API call failed: {response.status_code} - {response.text}"
 )


defnormalize_docstring(docstring):
"""Normalize docstring to strip leading/trailing spaces, remove markdown, and preserve multi-line format."""
ifnotdocstring:
return""

# Extract everything between triple quotes
# This assumes the docstring is surrounded by triple quotes
start=docstring.find('"""') +3# Skip past the first triple quotes
end=docstring.rfind('"""') # Find the last triple quotes

# Extract the content between the triple quotes
ifstart>2andend>start:
docstring=docstring[start:end].strip() # Remove the triple quotes and extra spaces

returndocstring


defget_functions_and_classes_without_docstrings(tree):
"""Parse the file and return functions and classes without docstrings, 
 considering top-level functions and functions inside a class only."""

functions_and_classes_without_docstrings= []

# We need to track the current class level to identify nested functions
class_stack= []

fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
# Determine the level of indentation
indentation_level=node.col_offset

# Debugging: Print function name, indentation level, and docstring presence
docstring=ast.get_docstring(node)
docstring_status="No docstring"ifdocstringisNoneelse"Has docstring"

# Skip functions with indentation level of 8 (nested functions inside other functions)
ifindentation_level==8:
# print(f"Skipping nested function: {node.name}")
continue
else:
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('function', node))

elifisinstance(node, ast.ClassDef):
# Add the class to the class stack, as we're inside a class now
class_stack.append(node)

# Check if the class has a docstring
docstring=ast.get_docstring(node)
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('class', node))

# Pop the class stack when leaving a class body
ifisinstance(node, ast.ClassDef):
class_stack.pop()

returnfunctions_and_classes_without_docstrings


defgenerate_docstring(object_code, openai_provider, cache, object_name, object_type, filename):
"""Generate a docstring using OpenAIProvider."""
code_digest=hashlib.sha256(object_code.encode()).hexdigest()
entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()
ifentryincacheandcode_digest==cache[entry]['last_hash']:
print(f"Using cached docstring for {object_type}: {object_name}")
returncache[entry]['docstring'] # Return cached docstring if it exists

ifentryincacheandcode_digest!=cache[entry]['last_hash']:
print(f"Has for {object_type}: {object_name} changed. Getting new docstring.")

prompt=f"""
 You will be provided with a Python {object_type} (function or class),
 including its code. Based on this, your task is to generate a Python docstring for it.
 Use the function signature and body to infer the purpose of the function/class,
 the arguments it takes, the return value, and any exceptions it may raise.

 Follow these guidelines to create the docstring:

 1. Summary: Provide a concise summary of the function/class's purpose.
 Focus on what the function does and its main goal.
 2. Arguments: List the parameters, their types, and a brief description
 of what each one represents.
 3. Return: If the function has a return value, describe the return type
 and what it represents. If there's no return, OMIT THE SECTION.
 4. Exceptions: If the function raises any exceptions, list them with descriptions.
 If no exceptions are raised, OMIT THE SECTION.
 5. Side Effects (if applicable): If the function has side effects
 (e.g., modifies global state, interacts with external services), mention them.
 OMIT THE SECTION if it is not lear in the code.
 6. Algorithm or Key Logic (optional): If the function is complex,
 provide a high-level outline of the logic or algorithm involved.
 OMIT THE SECTION if it is not lear in the code.

 Ensure the docstring follows Python's standard docstring conventions and provides
 just enough detail to make the function understandable and usable without overwhelming the reader.
 """

# Get the docstring from OpenAI
docstring=openai_provider.improve_text(prompt, object_code)

# Normalize the docstring by stripping leading/trailing whitespace
normalized_docstring=normalize_docstring(docstring)

# If the docstring is empty, set a fallback docstring
ifnotnormalized_docstring:
print(f"Warning: Generated docstring for {object_type} '{object_name}' is empty. Using a fallback docstring.")
normalized_docstring=f"TODO: Add a docstring for the {object_type} '{object_name}'."

tree=ast.parse(object_code)
indentation_level=0
fornodeinast.walk(tree):
ifisinstance(node, (ast.ClassDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+4# Add 4 spaces for method body
break
ifisinstance(node, (ast.FunctionDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+8# Add 4 spaces for method body
break

# Split the normalized docstring into lines and prefix each with the correct indentation
docstring_lines=normalized_docstring.splitlines()
indented_docstring="\n".join([f"{' '*indentation_level}{line}"forlineindocstring_lines])

# Cache the docstring
cache[entry] = {}
cache[entry]['docstring'] =indented_docstring
cache[entry]['last_hash'] =hashlib.sha256(object_code.encode()).hexdigest()

returnindented_docstring

defvalidate_syntax(code):
"""Validate Python code syntax by writing to a temporary file and checking it."""
try:
# Create a temporary file to check the syntax
withtempfile.NamedTemporaryFile(delete=False, mode='w', suffix=".py") astemp_file:
temp_file.write(code)
temp_filename=temp_file.name

# Run python -m py_compile to check syntax
result=subprocess.run(
 ['python', '-m', 'py_compile', temp_filename],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
 )

ifresult.returncode!=0:
print(code)

os.remove(temp_filename) # Remove the temporary file after checking

ifresult.returncode!=0:
print(f"Syntax error: {result.stderr.decode()}")
returnFalse
returnTrue

exceptExceptionase:
print(f"Error validating syntax: {e}")
returnFalse


defprocess_directory(directory, openai_provider, cache, cache_file):
"""Process all Python files in a directory and its subdirectories."""
forroot, dirs, filesinos.walk(directory):
forfileinfiles:
iffile.endswith(".py"): # Process only Python files
filepath=os.path.join(root, file)
process_file(filepath, openai_provider, cache, cache_file)


defload_cache(cache_file):
"""Load the cache from a file."""
ifos.path.exists(cache_file):
withopen(cache_file, 'r') asf:
returnjson.load(f)
return {}


defsave_cache(cache, cache_file):
"""Save the cache to a file."""
withopen(cache_file, 'w') asf:
json.dump(cache, f, indent=4)

defget_next(tree):
"""Get the next function or class without a docstring."""
# Assuming this function returns the next function/class that lacks a docstring.
# Return None when all functions/classes are documented
objects_without_docstrings=get_functions_and_classes_without_docstrings(tree)
ifobjects_without_docstrings:
returnobjects_without_docstrings[0] # Return the next object
returnNone# No more objects to process

defget_object_node(tree, object_name):
"""
 Traverses the AST to find and return the node for the object (function or class)
 with the specified name.

 Args:
 tree (ast.Module): The AST module to search through.
 object_name (str): The name of the function or class to find.

 Returns:
 ast.FunctionDef | ast.ClassDef: The AST node for the function or class.
 """
fornodeinast.walk(tree):
# Check for function or class definitions by name
ifisinstance(node, (ast.FunctionDef, ast.ClassDef)) andnode.name==object_name:
returnnode
returnNone

defprocess_file(filename, openai_provider, cache, cache_file):
"""Process the file: find functions, classes, generate docstrings, and save the result."""
print(f"Processing file: {filename}")

# Reread the file to ensure we are working with the latest version of the code
withopen(filename, 'r') asfile:
code=file.read()

# Start an infinite loop that will keep processing until there are no more items
whileTrue:
tree=ast.parse(code)
# Get the next function or class that doesn't have a docstring
next_object=get_next(tree)
ifnotnext_object:
print(f"All functions and classes have docstrings in {filename}.")
break# Exit the loop if there are no more functions or classes to document

object_type, object_node=next_object
# Get the object code (e.g., function or class)
object_code=ast.unparse(object_node)
object_name=object_node.name
print(f"Generating docstring for {object_type}: {object_name} in {filename}")

# Check if the function or class already has a docstring
ifast.get_docstring(object_node):
print(f"Skipping {object_type}: {object_name}, it already has a docstring.")
continue

# Generate the docstring using OpenAI or another provider
docstring=generate_docstring(object_code, openai_provider, cache, object_name, object_type, filename)

ifdocstring:
code=insert_docstring_to_object(code, object_node, docstring, object_type)

# Re-read the object code after the docstring insertion
updated_tree=ast.parse(code)
updated_object_node=get_object_node(updated_tree, object_name)
updated_object_code=ast.unparse(updated_object_node)

entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()

# Save the SHA-1 hash of the updated object_code into the cache
cache[entry]['last_hash'] =hashlib.sha1(updated_object_code.encode()).hexdigest()

# Save cache after processing each object
save_cache(cache, cache_file)
print(f"Docstrings added successfully to {object_type}: {object_name} in {filename}.")
else:
print(f"Something wrong with the docstring generation.")

# Validate Python syntax before saving
ifvalidate_syntax(code):
# Write the updated code back to the file
withopen(filename, 'w') asfile: 
file.write(code) 
else:
print(f"Syntax error in {filename}, skipping save.")

deffix_indentation(code):
"""Fix the indentation of the function signature and the code after the docstring."""

# Split the code into lines
code_lines=code.splitlines()

# Find the index where the function signature ends (the first occurrence of ':')
def_index=next(ifori, lineinenumerate(code_lines) ifline.strip().startswith('def '))

# Add 4 spaces to the function signature
code_lines[def_index] =' '+code_lines[def_index].lstrip()

# Now, we need to add indentation for all lines after the docstring
in_docstring=False
foriinrange(def_index+1, len(code_lines)):
line=code_lines[i]

# Detect the start and end of the docstring
ifline.strip().startswith('"""'):
ifin_docstring:
# End of the docstring
in_docstring=False
else:
# Start of the docstring
in_docstring=True

# Indent the line after the docstring
ifin_docstringisFalse:
code_lines[i] =' '+line# Add 4 spaces to the body of the function

# Rejoin the lines into the final code
updated_code="\n".join(code_lines)

returnupdated_code


defremove_duplicate_functions(code):
"""Remove duplicate function definitions based on the function signature."""
# Parse the code into an AST
tree=None
try:
tree=ast.parse(code)
exceptExceptionase:
print(code)
returncode

# Track function signatures
function_signatures=set()
functions_to_remove= []

# Walk through the AST and find function definitions
fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
function_signature=ast.unparse(node)

# If we have seen this function signature before, mark it for removal
iffunction_signatureinfunction_signatures:
functions_to_remove.append(node)
else:
function_signatures.add(function_signature)

# Remove duplicate functions from the code
forfunc_nodeinfunctions_to_remove:
# Removing the function by filtering out the lines corresponding to that function
code=remove_function_from_code(code, func_node)

returncode


defremove_function_from_code(code, function_node):
"""Remove the function from the code."""
# Split the code into lines
code_lines=code.splitlines()

# Get the start and end line numbers of the function
start_line=function_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(ast.unparse(function_node).splitlines())

# Remove the lines of the function from the code
return"\n".join(code_lines[:start_line] +code_lines[end_line+1:])


defextract_main_block(code):
"""Extracts everything from 'if __name__ == "__main__":' to the end of the file."""
lines=code.splitlines() # Split the code into lines
main_start_index=-1# Initialize the index

# Find the line number where '__name__' starts
fori, lineinenumerate(lines):
ifline.strip().startswith('if __name__ == "__main__":'):
main_start_index=i
break

ifmain_start_index==-1:
# print("No '__name__ == \"__main__\"' found in the code.")
return""

# Extract all lines starting from the '__name__' line to the end
return"\n".join(lines[main_start_index:])


definsert_docstring_to_object(code, object_node, docstring, object_type):
main=extract_main_block(code)
"""Insert the generated docstring into the function/class in the correct place."""
object_code=ast.unparse(object_node)

# Get the indentation level of the function or class signature using col_offset
indentation_level=' '* (object_node.col_offset+4) # Add 4 spaces for the body

# Check if the function or class already has a docstring, if so, remove it
existing_docstring=ast.get_docstring(object_node)
ifexisting_docstring:
# Remove the existing docstring from the object_code
object_code=object_code.replace(f'"""{existing_docstring}"""', '')

# The docstring is already properly formatted and indented, so we use it directly
docstring_code=f'{indentation_level}"""\n{docstring}\n{indentation_level}"""'

# Find the first occurrence of ':' to determine where the signature ends
signature_end=object_code.find(':')

func_start_lines=object_code[:signature_end+1]
func_end_lines=object_code[signature_end+1:]

# Insert the docstring immediately after the signature (before the body of the function/class)
object_code_with_docstring=func_start_lines+f'\n{docstring_code}\n'+func_end_lines

ifobject_node.col_offset>0:
object_code_with_docstring=fix_indentation(object_code_with_docstring)

# Get the start and end line numbers of the object_node
start_line=object_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(object_code.splitlines()) -1

# Split the original code into lines
code_lines=code.splitlines()

print("================================")
print(pprint.pprint(code_lines[:start_line]))
print("================================")

# Replace the lines between start_line and end_line with the new code containing the docstring
updated_code_lines= (
code_lines[:start_line] +
 [object_code_with_docstring] +
code_lines[end_line+1:]
 )

# Join the lines back into a single string
updated_code="\n".join(updated_code_lines)
updated_code=remove_duplicate_functions(updated_code)
main_after_update=extract_main_block(updated_code)

ifmain!=main_after_update:
updated_code=updated_code+"\n\n"+main


# Ensure a line break after the inserted docstring for readability
ifnotupdated_code.endswith('\n'):
updated_code+='\n'

returnupdated_code


if__name__=='__main__':
# Get the directory or file from the command line argument
iflen(sys.argv) !=2:
print("Usage: python add_docstrings.py <python_file_or_directory>")
sys.exit(1)

path=sys.argv[1]

# Initialize the OpenAIProvider
openai_provider=OpenAIProvider()

# Load the cache from the cache file (if it exists)
cache_file="docstring_cache.json"
cache=load_cache(cache_file)

ifos.path.isdir(path): # If it's a directory, process it recursively
process_directory(path, openai_provider, cache, cache_file)
elifos.path.isfile(path): # If it's a single file, process it
process_file(path, openai_provider, cache, cache_file)
else:
print(f"The provided path '{path}' is not valid.")
sys.exit(1)

# Save the updated cache to the file
save_cache(cache, cache_file)
#!/usr/bin/python
importast
importos
importrequests
importsys
importjson
importtempfile
importsubprocess
importhashlib
importpprint
importjson


classOpenAIProvider:
def__init__(self):
self.api_key=os.getenv("AI_API_KEY")
ifnotself.api_key:
raiseEnvironmentError("AI_API_KEY not set in environment.")
self.endpoint="https://api.openai.com/v1/chat/completions"
self.model=os.getenv("OPENAI_MODEL", "gpt-4o-mini")

defimprove_text(self, prompt: str, text: str) ->str:
headers= {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
 }

body= {
"model": self.model,
"messages": [
 {"role": "system", "content": prompt},
 {"role": "user", "content": text},
 ],
"temperature": 0.3,
 }

response=requests.post(self.endpoint, json=body, headers=headers, timeout=60)
ifresponse.status_code==200:
returnresponse.json()["choices"][0]["message"]["content"].strip()

raiseException(
f"OpenAI API call failed: {response.status_code} - {response.text}"
 )


defnormalize_docstring(docstring):
"""Normalize docstring to strip leading/trailing spaces, remove markdown, and preserve multi-line format."""
ifnotdocstring:
return""

# Extract everything between triple quotes
# This assumes the docstring is surrounded by triple quotes
start=docstring.find('"""') +3# Skip past the first triple quotes
end=docstring.rfind('"""') # Find the last triple quotes

# Extract the content between the triple quotes
ifstart>2andend>start:
docstring=docstring[start:end].strip() # Remove the triple quotes and extra spaces

returndocstring


defget_functions_and_classes_without_docstrings(tree):
"""Parse the file and return functions and classes without docstrings, 
 considering top-level functions and functions inside a class only."""

functions_and_classes_without_docstrings= []

# We need to track the current class level to identify nested functions
class_stack= []

fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
# Determine the level of indentation
indentation_level=node.col_offset

# Debugging: Print function name, indentation level, and docstring presence
docstring=ast.get_docstring(node)
docstring_status="No docstring"ifdocstringisNoneelse"Has docstring"

# Skip functions with indentation level of 8 (nested functions inside other functions)
ifindentation_level==8:
# print(f"Skipping nested function: {node.name}")
continue
else:
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('function', node))

elifisinstance(node, ast.ClassDef):
# Add the class to the class stack, as we're inside a class now
class_stack.append(node)

# Check if the class has a docstring
docstring=ast.get_docstring(node)
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('class', node))

# Pop the class stack when leaving a class body
ifisinstance(node, ast.ClassDef):
class_stack.pop()

returnfunctions_and_classes_without_docstrings


defgenerate_docstring(object_code, openai_provider, cache, object_name, object_type, filename):
"""Generate a docstring using OpenAIProvider."""
code_digest=hashlib.sha256(object_code.encode()).hexdigest()
entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()
ifentryincacheandcode_digest==cache[entry]['last_hash']:
print(f"Using cached docstring for {object_type}: {object_name}")
returncache[entry]['docstring'] # Return cached docstring if it exists

ifentryincacheandcode_digest!=cache[entry]['last_hash']:
print(f"Has for {object_type}: {object_name} changed. Getting new docstring.")

prompt=f"""
 You will be provided with a Python {object_type} (function or class),
 including its code. Based on this, your task is to generate a Python docstring for it.
 Use the function signature and body to infer the purpose of the function/class,
 the arguments it takes, the return value, and any exceptions it may raise.

 Follow these guidelines to create the docstring:

 1. Summary: Provide a concise summary of the function/class's purpose.
 Focus on what the function does and its main goal.
 2. Arguments: List the parameters, their types, and a brief description
 of what each one represents.
 3. Return: If the function has a return value, describe the return type
 and what it represents. If there's no return, OMIT THE SECTION.
 4. Exceptions: If the function raises any exceptions, list them with descriptions.
 If no exceptions are raised, OMIT THE SECTION.
 5. Side Effects (if applicable): If the function has side effects
 (e.g., modifies global state, interacts with external services), mention them.
 OMIT THE SECTION if it is not lear in the code.
 6. Algorithm or Key Logic (optional): If the function is complex,
 provide a high-level outline of the logic or algorithm involved.
 OMIT THE SECTION if it is not lear in the code.

 Ensure the docstring follows Python's standard docstring conventions and provides
 just enough detail to make the function understandable and usable without overwhelming the reader.
 """

# Get the docstring from OpenAI
docstring=openai_provider.improve_text(prompt, object_code)

# Normalize the docstring by stripping leading/trailing whitespace
normalized_docstring=normalize_docstring(docstring)

# If the docstring is empty, set a fallback docstring
ifnotnormalized_docstring:
print(f"Warning: Generated docstring for {object_type} '{object_name}' is empty. Using a fallback docstring.")
normalized_docstring=f"TODO: Add a docstring for the {object_type} '{object_name}'."

tree=ast.parse(object_code)
indentation_level=0
fornodeinast.walk(tree):
ifisinstance(node, (ast.ClassDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+4# Add 4 spaces for method body
break
ifisinstance(node, (ast.FunctionDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+8# Add 4 spaces for method body
break

# Split the normalized docstring into lines and prefix each with the correct indentation
docstring_lines=normalized_docstring.splitlines()
indented_docstring="\n".join([f"{' '*indentation_level}{line}"forlineindocstring_lines])

# Cache the docstring
cache[entry] = {}
cache[entry]['docstring'] =indented_docstring
cache[entry]['last_hash'] =hashlib.sha256(object_code.encode()).hexdigest()

returnindented_docstring

defvalidate_syntax(code):
"""Validate Python code syntax by writing to a temporary file and checking it."""
try:
# Create a temporary file to check the syntax
withtempfile.NamedTemporaryFile(delete=False, mode='w', suffix=".py") astemp_file:
temp_file.write(code)
temp_filename=temp_file.name

# Run python -m py_compile to check syntax
result=subprocess.run(
 ['python', '-m', 'py_compile', temp_filename],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
 )

ifresult.returncode!=0:
print(code)

os.remove(temp_filename) # Remove the temporary file after checking

ifresult.returncode!=0:
print(f"Syntax error: {result.stderr.decode()}")
returnFalse
returnTrue

exceptExceptionase:
print(f"Error validating syntax: {e}")
returnFalse


defprocess_directory(directory, openai_provider, cache, cache_file):
"""Process all Python files in a directory and its subdirectories."""
forroot, dirs, filesinos.walk(directory):
forfileinfiles:
iffile.endswith(".py"): # Process only Python files
filepath=os.path.join(root, file)
process_file(filepath, openai_provider, cache, cache_file)


defload_cache(cache_file):
"""Load the cache from a file."""
ifos.path.exists(cache_file):
withopen(cache_file, 'r') asf:
returnjson.load(f)
return {}


defsave_cache(cache, cache_file):
"""Save the cache to a file."""
withopen(cache_file, 'w') asf:
json.dump(cache, f, indent=4)

defget_next(tree):
"""Get the next function or class without a docstring."""
# Assuming this function returns the next function/class that lacks a docstring.
# Return None when all functions/classes are documented
objects_without_docstrings=get_functions_and_classes_without_docstrings(tree)
ifobjects_without_docstrings:
returnobjects_without_docstrings[0] # Return the next object
returnNone# No more objects to process

defget_object_node(tree, object_name):
"""
 Traverses the AST to find and return the node for the object (function or class)
 with the specified name.

 Args:
 tree (ast.Module): The AST module to search through.
 object_name (str): The name of the function or class to find.

 Returns:
 ast.FunctionDef | ast.ClassDef: The AST node for the function or class.
 """
fornodeinast.walk(tree):
# Check for function or class definitions by name
ifisinstance(node, (ast.FunctionDef, ast.ClassDef)) andnode.name==object_name:
returnnode
returnNone

defprocess_file(filename, openai_provider, cache, cache_file):
"""Process the file: find functions, classes, generate docstrings, and save the result."""
print(f"Processing file: {filename}")

# Reread the file to ensure we are working with the latest version of the code
withopen(filename, 'r') asfile:
code=file.read()

# Start an infinite loop that will keep processing until there are no more items
whileTrue:
tree=ast.parse(code)
# Get the next function or class that doesn't have a docstring
next_object=get_next(tree)
ifnotnext_object:
print(f"All functions and classes have docstrings in {filename}.")
break# Exit the loop if there are no more functions or classes to document

object_type, object_node=next_object
# Get the object code (e.g., function or class)
object_code=ast.unparse(object_node)
object_name=object_node.name
print(f"Generating docstring for {object_type}: {object_name} in {filename}")

# Check if the function or class already has a docstring
ifast.get_docstring(object_node):
print(f"Skipping {object_type}: {object_name}, it already has a docstring.")
continue

# Generate the docstring using OpenAI or another provider
docstring=generate_docstring(object_code, openai_provider, cache, object_name, object_type, filename)

ifdocstring:
code=insert_docstring_to_object(code, object_node, docstring, object_type)

# Re-read the object code after the docstring insertion
updated_tree=ast.parse(code)
updated_object_node=get_object_node(updated_tree, object_name)
updated_object_code=ast.unparse(updated_object_node)

entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()

# Save the SHA-1 hash of the updated object_code into the cache
cache[entry]['last_hash'] =hashlib.sha1(updated_object_code.encode()).hexdigest()

# Save cache after processing each object
save_cache(cache, cache_file)
print(f"Docstrings added successfully to {object_type}: {object_name} in {filename}.")
else:
print(f"Something wrong with the docstring generation.")

# Validate Python syntax before saving
ifvalidate_syntax(code):
# Write the updated code back to the file
withopen(filename, 'w') asfile: 
file.write(code) 
else:
print(f"Syntax error in {filename}, skipping save.")

deffix_indentation(code):
"""Fix the indentation of the function signature and the code after the docstring."""

# Split the code into lines
code_lines=code.splitlines()

# Find the index where the function signature ends (the first occurrence of ':')
def_index=next(ifori, lineinenumerate(code_lines) ifline.strip().startswith('def '))

# Add 4 spaces to the function signature
code_lines[def_index] =' '+code_lines[def_index].lstrip()

# Now, we need to add indentation for all lines after the docstring
in_docstring=False
foriinrange(def_index+1, len(code_lines)):
line=code_lines[i]

# Detect the start and end of the docstring
ifline.strip().startswith('"""'):
ifin_docstring:
# End of the docstring
in_docstring=False
else:
# Start of the docstring
in_docstring=True

# Indent the line after the docstring
ifin_docstringisFalse:
code_lines[i] =' '+line# Add 4 spaces to the body of the function

# Rejoin the lines into the final code
updated_code="\n".join(code_lines)

returnupdated_code


defremove_duplicate_functions(code):
"""Remove duplicate function definitions based on the function signature."""
# Parse the code into an AST
tree=None
try:
tree=ast.parse(code)
exceptExceptionase:
print(code)
returncode

# Track function signatures
function_signatures=set()
functions_to_remove= []

# Walk through the AST and find function definitions
fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
function_signature=ast.unparse(node)

# If we have seen this function signature before, mark it for removal
iffunction_signatureinfunction_signatures:
functions_to_remove.append(node)
else:
function_signatures.add(function_signature)

# Remove duplicate functions from the code
forfunc_nodeinfunctions_to_remove:
# Removing the function by filtering out the lines corresponding to that function
code=remove_function_from_code(code, func_node)

returncode


defremove_function_from_code(code, function_node):
"""Remove the function from the code."""
# Split the code into lines
code_lines=code.splitlines()

# Get the start and end line numbers of the function
start_line=function_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(ast.unparse(function_node).splitlines())

# Remove the lines of the function from the code
return"\n".join(code_lines[:start_line] +code_lines[end_line+1:])


defextract_main_block(code):
"""Extracts everything from 'if __name__ == "__main__":' to the end of the file."""
lines=code.splitlines() # Split the code into lines
main_start_index=-1# Initialize the index

# Find the line number where '__name__' starts
fori, lineinenumerate(lines):
ifline.strip().startswith('if __name__ == "__main__":'):
main_start_index=i
break

ifmain_start_index==-1:
# print("No '__name__ == \"__main__\"' found in the code.")
return""

# Extract all lines starting from the '__name__' line to the end
return"\n".join(lines[main_start_index:])


definsert_docstring_to_object(code, object_node, docstring, object_type):
main=extract_main_block(code)
"""Insert the generated docstring into the function/class in the correct place."""
object_code=ast.unparse(object_node)

# Get the indentation level of the function or class signature using col_offset
indentation_level=' '* (object_node.col_offset+4) # Add 4 spaces for the body

# Check if the function or class already has a docstring, if so, remove it
existing_docstring=ast.get_docstring(object_node)
ifexisting_docstring:
# Remove the existing docstring from the object_code
object_code=object_code.replace(f'"""{existing_docstring}"""', '')

# The docstring is already properly formatted and indented, so we use it directly
docstring_code=f'{indentation_level}"""\n{docstring}\n{indentation_level}"""'

# Find the first occurrence of ':' to determine where the signature ends
signature_end=object_code.find(':')

func_start_lines=object_code[:signature_end+1]
func_end_lines=object_code[signature_end+1:]

# Insert the docstring immediately after the signature (before the body of the function/class)
object_code_with_docstring=func_start_lines+f'\n{docstring_code}\n'+func_end_lines

ifobject_node.col_offset>0:
object_code_with_docstring=fix_indentation(object_code_with_docstring)

# Get the start and end line numbers of the object_node
start_line=object_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(object_code.splitlines()) -1

# Split the original code into lines
code_lines=code.splitlines()

print("================================")
print(pprint.pprint(code_lines[:start_line]))
print("================================")

# Replace the lines between start_line and end_line with the new code containing the docstring
updated_code_lines= (
code_lines[:start_line] +
 [object_code_with_docstring] +
code_lines[end_line+1:]
 )

# Join the lines back into a single string
updated_code="\n".join(updated_code_lines)
updated_code=remove_duplicate_functions(updated_code)
main_after_update=extract_main_block(updated_code)

ifmain!=main_after_update:
updated_code=updated_code+"\n\n"+main


# Ensure a line break after the inserted docstring for readability
ifnotupdated_code.endswith('\n'):
updated_code+='\n'

returnupdated_code


if__name__=='__main__':
# Get the directory or file from the command line argument
iflen(sys.argv) !=2:
print("Usage: python add_docstrings.py <python_file_or_directory>")
sys.exit(1)

path=sys.argv[1]

# Initialize the OpenAIProvider
openai_provider=OpenAIProvider()

# Load the cache from the cache file (if it exists)
cache_file="docstring_cache.json"
cache=load_cache(cache_file)

ifos.path.isdir(path): # If it's a directory, process it recursively
process_directory(path, openai_provider, cache, cache_file)
elifos.path.isfile(path): # If it's a single file, process it
process_file(path, openai_provider, cache, cache_file)
else:
print(f"The provided path '{path}' is not valid.")
sys.exit(1)

# Save the updated cache to the file
save_cache(cache, cache_file)
#!/usr/bin/python
importast
importos
importrequests
importsys
importjson
importtempfile
importsubprocess
importhashlib
importpprint
importjson


classOpenAIProvider:
def__init__(self):
self.api_key=os.getenv("AI_API_KEY")
ifnotself.api_key:
raiseEnvironmentError("AI_API_KEY not set in environment.")
self.endpoint="https://api.openai.com/v1/chat/completions"
self.model=os.getenv("OPENAI_MODEL", "gpt-4o-mini")

defimprove_text(self, prompt: str, text: str) ->str:
headers= {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
 }

body= {
"model": self.model,
"messages": [
 {"role": "system", "content": prompt},
 {"role": "user", "content": text},
 ],
"temperature": 0.3,
 }

response=requests.post(self.endpoint, json=body, headers=headers, timeout=60)
ifresponse.status_code==200:
returnresponse.json()["choices"][0]["message"]["content"].strip()

raiseException(
f"OpenAI API call failed: {response.status_code} - {response.text}"
 )


defnormalize_docstring(docstring):
"""Normalize docstring to strip leading/trailing spaces, remove markdown, and preserve multi-line format."""
ifnotdocstring:
return""

# Extract everything between triple quotes
# This assumes the docstring is surrounded by triple quotes
start=docstring.find('"""') +3# Skip past the first triple quotes
end=docstring.rfind('"""') # Find the last triple quotes

# Extract the content between the triple quotes
ifstart>2andend>start:
docstring=docstring[start:end].strip() # Remove the triple quotes and extra spaces

returndocstring


defget_functions_and_classes_without_docstrings(tree):
"""Parse the file and return functions and classes without docstrings, 
 considering top-level functions and functions inside a class only."""

functions_and_classes_without_docstrings= []

# We need to track the current class level to identify nested functions
class_stack= []

fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
# Determine the level of indentation
indentation_level=node.col_offset

# Debugging: Print function name, indentation level, and docstring presence
docstring=ast.get_docstring(node)
docstring_status="No docstring"ifdocstringisNoneelse"Has docstring"

# Skip functions with indentation level of 8 (nested functions inside other functions)
ifindentation_level==8:
# print(f"Skipping nested function: {node.name}")
continue
else:
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('function', node))

elifisinstance(node, ast.ClassDef):
# Add the class to the class stack, as we're inside a class now
class_stack.append(node)

# Check if the class has a docstring
docstring=ast.get_docstring(node)
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('class', node))

# Pop the class stack when leaving a class body
ifisinstance(node, ast.ClassDef):
class_stack.pop()

returnfunctions_and_classes_without_docstrings


defgenerate_docstring(object_code, openai_provider, cache, object_name, object_type, filename):
"""Generate a docstring using OpenAIProvider."""
code_digest=hashlib.sha256(object_code.encode()).hexdigest()
entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()
ifentryincacheandcode_digest==cache[entry]['last_hash']:
print(f"Using cached docstring for {object_type}: {object_name}")
returncache[entry]['docstring'] # Return cached docstring if it exists

ifentryincacheandcode_digest!=cache[entry]['last_hash']:
print(f"Has for {object_type}: {object_name} changed. Getting new docstring.")

prompt=f"""
 You will be provided with a Python {object_type} (function or class),
 including its code. Based on this, your task is to generate a Python docstring for it.
 Use the function signature and body to infer the purpose of the function/class,
 the arguments it takes, the return value, and any exceptions it may raise.

 Follow these guidelines to create the docstring:

 1. Summary: Provide a concise summary of the function/class's purpose.
 Focus on what the function does and its main goal.
 2. Arguments: List the parameters, their types, and a brief description
 of what each one represents.
 3. Return: If the function has a return value, describe the return type
 and what it represents. If there's no return, OMIT THE SECTION.
 4. Exceptions: If the function raises any exceptions, list them with descriptions.
 If no exceptions are raised, OMIT THE SECTION.
 5. Side Effects (if applicable): If the function has side effects
 (e.g., modifies global state, interacts with external services), mention them.
 OMIT THE SECTION if it is not lear in the code.
 6. Algorithm or Key Logic (optional): If the function is complex,
 provide a high-level outline of the logic or algorithm involved.
 OMIT THE SECTION if it is not lear in the code.

 Ensure the docstring follows Python's standard docstring conventions and provides
 just enough detail to make the function understandable and usable without overwhelming the reader.
 """

# Get the docstring from OpenAI
docstring=openai_provider.improve_text(prompt, object_code)

# Normalize the docstring by stripping leading/trailing whitespace
normalized_docstring=normalize_docstring(docstring)

# If the docstring is empty, set a fallback docstring
ifnotnormalized_docstring:
print(f"Warning: Generated docstring for {object_type} '{object_name}' is empty. Using a fallback docstring.")
normalized_docstring=f"TODO: Add a docstring for the {object_type} '{object_name}'."

tree=ast.parse(object_code)
indentation_level=0
fornodeinast.walk(tree):
ifisinstance(node, (ast.ClassDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+4# Add 4 spaces for method body
break
ifisinstance(node, (ast.FunctionDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+8# Add 4 spaces for method body
break

# Split the normalized docstring into lines and prefix each with the correct indentation
docstring_lines=normalized_docstring.splitlines()
indented_docstring="\n".join([f"{' '*indentation_level}{line}"forlineindocstring_lines])

# Cache the docstring
cache[entry] = {}
cache[entry]['docstring'] =indented_docstring
cache[entry]['last_hash'] =hashlib.sha256(object_code.encode()).hexdigest()

returnindented_docstring

defvalidate_syntax(code):
"""Validate Python code syntax by writing to a temporary file and checking it."""
try:
# Create a temporary file to check the syntax
withtempfile.NamedTemporaryFile(delete=False, mode='w', suffix=".py") astemp_file:
temp_file.write(code)
temp_filename=temp_file.name

# Run python -m py_compile to check syntax
result=subprocess.run(
 ['python', '-m', 'py_compile', temp_filename],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
 )

ifresult.returncode!=0:
print(code)

os.remove(temp_filename) # Remove the temporary file after checking

ifresult.returncode!=0:
print(f"Syntax error: {result.stderr.decode()}")
returnFalse
returnTrue

exceptExceptionase:
print(f"Error validating syntax: {e}")
returnFalse


defprocess_directory(directory, openai_provider, cache, cache_file):
"""Process all Python files in a directory and its subdirectories."""
forroot, dirs, filesinos.walk(directory):
forfileinfiles:
iffile.endswith(".py"): # Process only Python files
filepath=os.path.join(root, file)
process_file(filepath, openai_provider, cache, cache_file)


defload_cache(cache_file):
"""Load the cache from a file."""
ifos.path.exists(cache_file):
withopen(cache_file, 'r') asf:
returnjson.load(f)
return {}


defsave_cache(cache, cache_file):
"""Save the cache to a file."""
withopen(cache_file, 'w') asf:
json.dump(cache, f, indent=4)

defget_next(tree):
"""Get the next function or class without a docstring."""
# Assuming this function returns the next function/class that lacks a docstring.
# Return None when all functions/classes are documented
objects_without_docstrings=get_functions_and_classes_without_docstrings(tree)
ifobjects_without_docstrings:
returnobjects_without_docstrings[0] # Return the next object
returnNone# No more objects to process

defget_object_node(tree, object_name):
"""
 Traverses the AST to find and return the node for the object (function or class)
 with the specified name.

 Args:
 tree (ast.Module): The AST module to search through.
 object_name (str): The name of the function or class to find.

 Returns:
 ast.FunctionDef | ast.ClassDef: The AST node for the function or class.
 """
fornodeinast.walk(tree):
# Check for function or class definitions by name
ifisinstance(node, (ast.FunctionDef, ast.ClassDef)) andnode.name==object_name:
returnnode
returnNone

defprocess_file(filename, openai_provider, cache, cache_file):
"""Process the file: find functions, classes, generate docstrings, and save the result."""
print(f"Processing file: {filename}")

# Reread the file to ensure we are working with the latest version of the code
withopen(filename, 'r') asfile:
code=file.read()

# Start an infinite loop that will keep processing until there are no more items
whileTrue:
tree=ast.parse(code)
# Get the next function or class that doesn't have a docstring
next_object=get_next(tree)
ifnotnext_object:
print(f"All functions and classes have docstrings in {filename}.")
break# Exit the loop if there are no more functions or classes to document

object_type, object_node=next_object
# Get the object code (e.g., function or class)
object_code=ast.unparse(object_node)
object_name=object_node.name
print(f"Generating docstring for {object_type}: {object_name} in {filename}")

# Check if the function or class already has a docstring
ifast.get_docstring(object_node):
print(f"Skipping {object_type}: {object_name}, it already has a docstring.")
continue

# Generate the docstring using OpenAI or another provider
docstring=generate_docstring(object_code, openai_provider, cache, object_name, object_type, filename)

ifdocstring:
code=insert_docstring_to_object(code, object_node, docstring, object_type)

# Re-read the object code after the docstring insertion
updated_tree=ast.parse(code)
updated_object_node=get_object_node(updated_tree, object_name)
updated_object_code=ast.unparse(updated_object_node)

entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()

# Save the SHA-1 hash of the updated object_code into the cache
cache[entry]['last_hash'] =hashlib.sha1(updated_object_code.encode()).hexdigest()

# Save cache after processing each object
save_cache(cache, cache_file)
print(f"Docstrings added successfully to {object_type}: {object_name} in {filename}.")
else:
print(f"Something wrong with the docstring generation.")

# Validate Python syntax before saving
ifvalidate_syntax(code):
# Write the updated code back to the file
withopen(filename, 'w') asfile: 
file.write(code) 
else:
print(f"Syntax error in {filename}, skipping save.")

deffix_indentation(code):
"""Fix the indentation of the function signature and the code after the docstring."""

# Split the code into lines
code_lines=code.splitlines()

# Find the index where the function signature ends (the first occurrence of ':')
def_index=next(ifori, lineinenumerate(code_lines) ifline.strip().startswith('def '))

# Add 4 spaces to the function signature
code_lines[def_index] =' '+code_lines[def_index].lstrip()

# Now, we need to add indentation for all lines after the docstring
in_docstring=False
foriinrange(def_index+1, len(code_lines)):
line=code_lines[i]

# Detect the start and end of the docstring
ifline.strip().startswith('"""'):
ifin_docstring:
# End of the docstring
in_docstring=False
else:
# Start of the docstring
in_docstring=True

# Indent the line after the docstring
ifin_docstringisFalse:
code_lines[i] =' '+line# Add 4 spaces to the body of the function

# Rejoin the lines into the final code
updated_code="\n".join(code_lines)

returnupdated_code


defremove_duplicate_functions(code):
"""Remove duplicate function definitions based on the function signature."""
# Parse the code into an AST
tree=None
try:
tree=ast.parse(code)
exceptExceptionase:
print(code)
returncode

# Track function signatures
function_signatures=set()
functions_to_remove= []

# Walk through the AST and find function definitions
fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
function_signature=ast.unparse(node)

# If we have seen this function signature before, mark it for removal
iffunction_signatureinfunction_signatures:
functions_to_remove.append(node)
else:
function_signatures.add(function_signature)

# Remove duplicate functions from the code
forfunc_nodeinfunctions_to_remove:
# Removing the function by filtering out the lines corresponding to that function
code=remove_function_from_code(code, func_node)

returncode


defremove_function_from_code(code, function_node):
"""Remove the function from the code."""
# Split the code into lines
code_lines=code.splitlines()

# Get the start and end line numbers of the function
start_line=function_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(ast.unparse(function_node).splitlines())

# Remove the lines of the function from the code
return"\n".join(code_lines[:start_line] +code_lines[end_line+1:])


defextract_main_block(code):
"""Extracts everything from 'if __name__ == "__main__":' to the end of the file."""
lines=code.splitlines() # Split the code into lines
main_start_index=-1# Initialize the index

# Find the line number where '__name__' starts
fori, lineinenumerate(lines):
ifline.strip().startswith('if __name__ == "__main__":'):
main_start_index=i
break

ifmain_start_index==-1:
# print("No '__name__ == \"__main__\"' found in the code.")
return""

# Extract all lines starting from the '__name__' line to the end
return"\n".join(lines[main_start_index:])


definsert_docstring_to_object(code, object_node, docstring, object_type):
main=extract_main_block(code)
"""Insert the generated docstring into the function/class in the correct place."""
object_code=ast.unparse(object_node)

# Get the indentation level of the function or class signature using col_offset
indentation_level=' '* (object_node.col_offset+4) # Add 4 spaces for the body

# Check if the function or class already has a docstring, if so, remove it
existing_docstring=ast.get_docstring(object_node)
ifexisting_docstring:
# Remove the existing docstring from the object_code
object_code=object_code.replace(f'"""{existing_docstring}"""', '')

# The docstring is already properly formatted and indented, so we use it directly
docstring_code=f'{indentation_level}"""\n{docstring}\n{indentation_level}"""'

# Find the first occurrence of ':' to determine where the signature ends
signature_end=object_code.find(':')

func_start_lines=object_code[:signature_end+1]
func_end_lines=object_code[signature_end+1:]

# Insert the docstring immediately after the signature (before the body of the function/class)
object_code_with_docstring=func_start_lines+f'\n{docstring_code}\n'+func_end_lines

ifobject_node.col_offset>0:
object_code_with_docstring=fix_indentation(object_code_with_docstring)

# Get the start and end line numbers of the object_node
start_line=object_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(object_code.splitlines()) -1

# Split the original code into lines
code_lines=code.splitlines()

print("================================")
print(pprint.pprint(code_lines[:start_line]))
print("================================")

# Replace the lines between start_line and end_line with the new code containing the docstring
updated_code_lines= (
code_lines[:start_line] +
 [object_code_with_docstring] +
code_lines[end_line+1:]
 )

# Join the lines back into a single string
updated_code="\n".join(updated_code_lines)
updated_code=remove_duplicate_functions(updated_code)
main_after_update=extract_main_block(updated_code)

ifmain!=main_after_update:
updated_code=updated_code+"\n\n"+main


# Ensure a line break after the inserted docstring for readability
ifnotupdated_code.endswith('\n'):
updated_code+='\n'

returnupdated_code


if__name__=='__main__':
# Get the directory or file from the command line argument
iflen(sys.argv) !=2:
print("Usage: python add_docstrings.py <python_file_or_directory>")
sys.exit(1)

path=sys.argv[1]

# Initialize the OpenAIProvider
openai_provider=OpenAIProvider()

# Load the cache from the cache file (if it exists)
cache_file="docstring_cache.json"
cache=load_cache(cache_file)

ifos.path.isdir(path): # If it's a directory, process it recursively
process_directory(path, openai_provider, cache, cache_file)
elifos.path.isfile(path): # If it's a single file, process it
process_file(path, openai_provider, cache, cache_file)
else:
print(f"The provided path '{path}' is not valid.")
sys.exit(1)

# Save the updated cache to the file
save_cache(cache, cache_file)
#!/usr/bin/python
importast
importos
importrequests
importsys
importjson
importtempfile
importsubprocess
importhashlib
importpprint
importjson


classOpenAIProvider:
def__init__(self):
self.api_key=os.getenv("AI_API_KEY")
ifnotself.api_key:
raiseEnvironmentError("AI_API_KEY not set in environment.")
self.endpoint="https://api.openai.com/v1/chat/completions"
self.model=os.getenv("OPENAI_MODEL", "gpt-4o-mini")

defimprove_text(self, prompt: str, text: str) ->str:
headers= {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
 }

body= {
"model": self.model,
"messages": [
 {"role": "system", "content": prompt},
 {"role": "user", "content": text},
 ],
"temperature": 0.3,
 }

response=requests.post(self.endpoint, json=body, headers=headers, timeout=60)
ifresponse.status_code==200:
returnresponse.json()["choices"][0]["message"]["content"].strip()

raiseException(
f"OpenAI API call failed: {response.status_code} - {response.text}"
 )


defnormalize_docstring(docstring):
"""Normalize docstring to strip leading/trailing spaces, remove markdown, and preserve multi-line format."""
ifnotdocstring:
return""

# Extract everything between triple quotes
# This assumes the docstring is surrounded by triple quotes
start=docstring.find('"""') +3# Skip past the first triple quotes
end=docstring.rfind('"""') # Find the last triple quotes

# Extract the content between the triple quotes
ifstart>2andend>start:
docstring=docstring[start:end].strip() # Remove the triple quotes and extra spaces

returndocstring


defget_functions_and_classes_without_docstrings(tree):
"""Parse the file and return functions and classes without docstrings, 
 considering top-level functions and functions inside a class only."""

functions_and_classes_without_docstrings= []

# We need to track the current class level to identify nested functions
class_stack= []

fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
# Determine the level of indentation
indentation_level=node.col_offset

# Debugging: Print function name, indentation level, and docstring presence
docstring=ast.get_docstring(node)
docstring_status="No docstring"ifdocstringisNoneelse"Has docstring"

# Skip functions with indentation level of 8 (nested functions inside other functions)
ifindentation_level==8:
# print(f"Skipping nested function: {node.name}")
continue
else:
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('function', node))

elifisinstance(node, ast.ClassDef):
# Add the class to the class stack, as we're inside a class now
class_stack.append(node)

# Check if the class has a docstring
docstring=ast.get_docstring(node)
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('class', node))

# Pop the class stack when leaving a class body
ifisinstance(node, ast.ClassDef):
class_stack.pop()

returnfunctions_and_classes_without_docstrings


defgenerate_docstring(object_code, openai_provider, cache, object_name, object_type, filename):
"""Generate a docstring using OpenAIProvider."""
code_digest=hashlib.sha256(object_code.encode()).hexdigest()
entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()
ifentryincacheandcode_digest==cache[entry]['last_hash']:
print(f"Using cached docstring for {object_type}: {object_name}")
returncache[entry]['docstring'] # Return cached docstring if it exists

ifentryincacheandcode_digest!=cache[entry]['last_hash']:
print(f"Has for {object_type}: {object_name} changed. Getting new docstring.")

prompt=f"""
 You will be provided with a Python {object_type} (function or class),
 including its code. Based on this, your task is to generate a Python docstring for it.
 Use the function signature and body to infer the purpose of the function/class,
 the arguments it takes, the return value, and any exceptions it may raise.

 Follow these guidelines to create the docstring:

 1. Summary: Provide a concise summary of the function/class's purpose.
 Focus on what the function does and its main goal.
 2. Arguments: List the parameters, their types, and a brief description
 of what each one represents.
 3. Return: If the function has a return value, describe the return type
 and what it represents. If there's no return, OMIT THE SECTION.
 4. Exceptions: If the function raises any exceptions, list them with descriptions.
 If no exceptions are raised, OMIT THE SECTION.
 5. Side Effects (if applicable): If the function has side effects
 (e.g., modifies global state, interacts with external services), mention them.
 OMIT THE SECTION if it is not lear in the code.
 6. Algorithm or Key Logic (optional): If the function is complex,
 provide a high-level outline of the logic or algorithm involved.
 OMIT THE SECTION if it is not lear in the code.

 Ensure the docstring follows Python's standard docstring conventions and provides
 just enough detail to make the function understandable and usable without overwhelming the reader.
 """

# Get the docstring from OpenAI
docstring=openai_provider.improve_text(prompt, object_code)

# Normalize the docstring by stripping leading/trailing whitespace
normalized_docstring=normalize_docstring(docstring)

# If the docstring is empty, set a fallback docstring
ifnotnormalized_docstring:
print(f"Warning: Generated docstring for {object_type} '{object_name}' is empty. Using a fallback docstring.")
normalized_docstring=f"TODO: Add a docstring for the {object_type} '{object_name}'."

tree=ast.parse(object_code)
indentation_level=0
fornodeinast.walk(tree):
ifisinstance(node, (ast.ClassDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+4# Add 4 spaces for method body
break
ifisinstance(node, (ast.FunctionDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+8# Add 4 spaces for method body
break

# Split the normalized docstring into lines and prefix each with the correct indentation
docstring_lines=normalized_docstring.splitlines()
indented_docstring="\n".join([f"{' '*indentation_level}{line}"forlineindocstring_lines])

# Cache the docstring
cache[entry] = {}
cache[entry]['docstring'] =indented_docstring
cache[entry]['last_hash'] =hashlib.sha256(object_code.encode()).hexdigest()

returnindented_docstring

defvalidate_syntax(code):
"""Validate Python code syntax by writing to a temporary file and checking it."""
try:
# Create a temporary file to check the syntax
withtempfile.NamedTemporaryFile(delete=False, mode='w', suffix=".py") astemp_file:
temp_file.write(code)
temp_filename=temp_file.name

# Run python -m py_compile to check syntax
result=subprocess.run(
 ['python', '-m', 'py_compile', temp_filename],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
 )

ifresult.returncode!=0:
print(code)

os.remove(temp_filename) # Remove the temporary file after checking

ifresult.returncode!=0:
print(f"Syntax error: {result.stderr.decode()}")
returnFalse
returnTrue

exceptExceptionase:
print(f"Error validating syntax: {e}")
returnFalse


defprocess_directory(directory, openai_provider, cache, cache_file):
"""Process all Python files in a directory and its subdirectories."""
forroot, dirs, filesinos.walk(directory):
forfileinfiles:
iffile.endswith(".py"): # Process only Python files
filepath=os.path.join(root, file)
process_file(filepath, openai_provider, cache, cache_file)


defload_cache(cache_file):
"""Load the cache from a file."""
ifos.path.exists(cache_file):
withopen(cache_file, 'r') asf:
returnjson.load(f)
return {}


defsave_cache(cache, cache_file):
"""Save the cache to a file."""
withopen(cache_file, 'w') asf:
json.dump(cache, f, indent=4)

defget_next(tree):
"""Get the next function or class without a docstring."""
# Assuming this function returns the next function/class that lacks a docstring.
# Return None when all functions/classes are documented
objects_without_docstrings=get_functions_and_classes_without_docstrings(tree)
ifobjects_without_docstrings:
returnobjects_without_docstrings[0] # Return the next object
returnNone# No more objects to process

defget_object_node(tree, object_name):
"""
 Traverses the AST to find and return the node for the object (function or class)
 with the specified name.

 Args:
 tree (ast.Module): The AST module to search through.
 object_name (str): The name of the function or class to find.

 Returns:
 ast.FunctionDef | ast.ClassDef: The AST node for the function or class.
 """
fornodeinast.walk(tree):
# Check for function or class definitions by name
ifisinstance(node, (ast.FunctionDef, ast.ClassDef)) andnode.name==object_name:
returnnode
returnNone

defprocess_file(filename, openai_provider, cache, cache_file):
"""Process the file: find functions, classes, generate docstrings, and save the result."""
print(f"Processing file: {filename}")

# Reread the file to ensure we are working with the latest version of the code
withopen(filename, 'r') asfile:
code=file.read()

# Start an infinite loop that will keep processing until there are no more items
whileTrue:
tree=ast.parse(code)
# Get the next function or class that doesn't have a docstring
next_object=get_next(tree)
ifnotnext_object:
print(f"All functions and classes have docstrings in {filename}.")
break# Exit the loop if there are no more functions or classes to document

object_type, object_node=next_object
# Get the object code (e.g., function or class)
object_code=ast.unparse(object_node)
object_name=object_node.name
print(f"Generating docstring for {object_type}: {object_name} in {filename}")

# Check if the function or class already has a docstring
ifast.get_docstring(object_node):
print(f"Skipping {object_type}: {object_name}, it already has a docstring.")
continue

# Generate the docstring using OpenAI or another provider
docstring=generate_docstring(object_code, openai_provider, cache, object_name, object_type, filename)

ifdocstring:
code=insert_docstring_to_object(code, object_node, docstring, object_type)

# Re-read the object code after the docstring insertion
updated_tree=ast.parse(code)
updated_object_node=get_object_node(updated_tree, object_name)
updated_object_code=ast.unparse(updated_object_node)

entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()

# Save the SHA-1 hash of the updated object_code into the cache
cache[entry]['last_hash'] =hashlib.sha1(updated_object_code.encode()).hexdigest()

# Save cache after processing each object
save_cache(cache, cache_file)
print(f"Docstrings added successfully to {object_type}: {object_name} in {filename}.")
else:
print(f"Something wrong with the docstring generation.")

# Validate Python syntax before saving
ifvalidate_syntax(code):
# Write the updated code back to the file
withopen(filename, 'w') asfile: 
file.write(code) 
else:
print(f"Syntax error in {filename}, skipping save.")

deffix_indentation(code):
"""Fix the indentation of the function signature and the code after the docstring."""

# Split the code into lines
code_lines=code.splitlines()

# Find the index where the function signature ends (the first occurrence of ':')
def_index=next(ifori, lineinenumerate(code_lines) ifline.strip().startswith('def '))

# Add 4 spaces to the function signature
code_lines[def_index] =' '+code_lines[def_index].lstrip()

# Now, we need to add indentation for all lines after the docstring
in_docstring=False
foriinrange(def_index+1, len(code_lines)):
line=code_lines[i]

# Detect the start and end of the docstring
ifline.strip().startswith('"""'):
ifin_docstring:
# End of the docstring
in_docstring=False
else:
# Start of the docstring
in_docstring=True

# Indent the line after the docstring
ifin_docstringisFalse:
code_lines[i] =' '+line# Add 4 spaces to the body of the function

# Rejoin the lines into the final code
updated_code="\n".join(code_lines)

returnupdated_code


defremove_duplicate_functions(code):
"""Remove duplicate function definitions based on the function signature."""
# Parse the code into an AST
tree=None
try:
tree=ast.parse(code)
exceptExceptionase:
print(code)
returncode

# Track function signatures
function_signatures=set()
functions_to_remove= []

# Walk through the AST and find function definitions
fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
function_signature=ast.unparse(node)

# If we have seen this function signature before, mark it for removal
iffunction_signatureinfunction_signatures:
functions_to_remove.append(node)
else:
function_signatures.add(function_signature)

# Remove duplicate functions from the code
forfunc_nodeinfunctions_to_remove:
# Removing the function by filtering out the lines corresponding to that function
code=remove_function_from_code(code, func_node)

returncode


defremove_function_from_code(code, function_node):
"""Remove the function from the code."""
# Split the code into lines
code_lines=code.splitlines()

# Get the start and end line numbers of the function
start_line=function_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(ast.unparse(function_node).splitlines())

# Remove the lines of the function from the code
return"\n".join(code_lines[:start_line] +code_lines[end_line+1:])


defextract_main_block(code):
"""Extracts everything from 'if __name__ == "__main__":' to the end of the file."""
lines=code.splitlines() # Split the code into lines
main_start_index=-1# Initialize the index

# Find the line number where '__name__' starts
fori, lineinenumerate(lines):
ifline.strip().startswith('if __name__ == "__main__":'):
main_start_index=i
break

ifmain_start_index==-1:
# print("No '__name__ == \"__main__\"' found in the code.")
return""

# Extract all lines starting from the '__name__' line to the end
return"\n".join(lines[main_start_index:])


definsert_docstring_to_object(code, object_node, docstring, object_type):
main=extract_main_block(code)
"""Insert the generated docstring into the function/class in the correct place."""
object_code=ast.unparse(object_node)

# Get the indentation level of the function or class signature using col_offset
indentation_level=' '* (object_node.col_offset+4) # Add 4 spaces for the body

# Check if the function or class already has a docstring, if so, remove it
existing_docstring=ast.get_docstring(object_node)
ifexisting_docstring:
# Remove the existing docstring from the object_code
object_code=object_code.replace(f'"""{existing_docstring}"""', '')

# The docstring is already properly formatted and indented, so we use it directly
docstring_code=f'{indentation_level}"""\n{docstring}\n{indentation_level}"""'

# Find the first occurrence of ':' to determine where the signature ends
signature_end=object_code.find(':')

func_start_lines=object_code[:signature_end+1]
func_end_lines=object_code[signature_end+1:]

# Insert the docstring immediately after the signature (before the body of the function/class)
object_code_with_docstring=func_start_lines+f'\n{docstring_code}\n'+func_end_lines

ifobject_node.col_offset>0:
object_code_with_docstring=fix_indentation(object_code_with_docstring)

# Get the start and end line numbers of the object_node
start_line=object_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(object_code.splitlines()) -1

# Split the original code into lines
code_lines=code.splitlines()

print("================================")
print(pprint.pprint(code_lines[:start_line]))
print("================================")

# Replace the lines between start_line and end_line with the new code containing the docstring
updated_code_lines= (
code_lines[:start_line] +
 [object_code_with_docstring] +
code_lines[end_line+1:]
 )

# Join the lines back into a single string
updated_code="\n".join(updated_code_lines)
updated_code=remove_duplicate_functions(updated_code)
main_after_update=extract_main_block(updated_code)

ifmain!=main_after_update:
updated_code=updated_code+"\n\n"+main


# Ensure a line break after the inserted docstring for readability
ifnotupdated_code.endswith('\n'):
updated_code+='\n'

returnupdated_code


if__name__=='__main__':
# Get the directory or file from the command line argument
iflen(sys.argv) !=2:
print("Usage: python add_docstrings.py <python_file_or_directory>")
sys.exit(1)

path=sys.argv[1]

# Initialize the OpenAIProvider
openai_provider=OpenAIProvider()

# Load the cache from the cache file (if it exists)
cache_file="docstring_cache.json"
cache=load_cache(cache_file)

ifos.path.isdir(path): # If it's a directory, process it recursively
process_directory(path, openai_provider, cache, cache_file)
elifos.path.isfile(path): # If it's a single file, process it
process_file(path, openai_provider, cache, cache_file)
else:
print(f"The provided path '{path}' is not valid.")
sys.exit(1)

# Save the updated cache to the file
save_cache(cache, cache_file)
#!/usr/bin/python
importast
importos
importrequests
importsys
importjson
importtempfile
importsubprocess
importhashlib
importpprint
importjson


classOpenAIProvider:
def__init__(self):
self.api_key=os.getenv("AI_API_KEY")
ifnotself.api_key:
raiseEnvironmentError("AI_API_KEY not set in environment.")
self.endpoint="https://api.openai.com/v1/chat/completions"
self.model=os.getenv("OPENAI_MODEL", "gpt-4o-mini")

defimprove_text(self, prompt: str, text: str) ->str:
headers= {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
 }

body= {
"model": self.model,
"messages": [
 {"role": "system", "content": prompt},
 {"role": "user", "content": text},
 ],
"temperature": 0.3,
 }

response=requests.post(self.endpoint, json=body, headers=headers, timeout=60)
ifresponse.status_code==200:
returnresponse.json()["choices"][0]["message"]["content"].strip()

raiseException(
f"OpenAI API call failed: {response.status_code} - {response.text}"
 )


defnormalize_docstring(docstring):
"""Normalize docstring to strip leading/trailing spaces, remove markdown, and preserve multi-line format."""
ifnotdocstring:
return""

# Extract everything between triple quotes
# This assumes the docstring is surrounded by triple quotes
start=docstring.find('"""') +3# Skip past the first triple quotes
end=docstring.rfind('"""') # Find the last triple quotes

# Extract the content between the triple quotes
ifstart>2andend>start:
docstring=docstring[start:end].strip() # Remove the triple quotes and extra spaces

returndocstring


defget_functions_and_classes_without_docstrings(tree):
"""Parse the file and return functions and classes without docstrings, 
 considering top-level functions and functions inside a class only."""

functions_and_classes_without_docstrings= []

# We need to track the current class level to identify nested functions
class_stack= []

fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
# Determine the level of indentation
indentation_level=node.col_offset

# Debugging: Print function name, indentation level, and docstring presence
docstring=ast.get_docstring(node)
docstring_status="No docstring"ifdocstringisNoneelse"Has docstring"

# Skip functions with indentation level of 8 (nested functions inside other functions)
ifindentation_level==8:
# print(f"Skipping nested function: {node.name}")
continue
else:
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('function', node))

elifisinstance(node, ast.ClassDef):
# Add the class to the class stack, as we're inside a class now
class_stack.append(node)

# Check if the class has a docstring
docstring=ast.get_docstring(node)
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('class', node))

# Pop the class stack when leaving a class body
ifisinstance(node, ast.ClassDef):
class_stack.pop()

returnfunctions_and_classes_without_docstrings


defgenerate_docstring(object_code, openai_provider, cache, object_name, object_type, filename):
"""Generate a docstring using OpenAIProvider."""
code_digest=hashlib.sha256(object_code.encode()).hexdigest()
entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()
ifentryincacheandcode_digest==cache[entry]['last_hash']:
print(f"Using cached docstring for {object_type}: {object_name}")
returncache[entry]['docstring'] # Return cached docstring if it exists

ifentryincacheandcode_digest!=cache[entry]['last_hash']:
print(f"Has for {object_type}: {object_name} changed. Getting new docstring.")

prompt=f"""
 You will be provided with a Python {object_type} (function or class),
 including its code. Based on this, your task is to generate a Python docstring for it.
 Use the function signature and body to infer the purpose of the function/class,
 the arguments it takes, the return value, and any exceptions it may raise.

 Follow these guidelines to create the docstring:

 1. Summary: Provide a concise summary of the function/class's purpose.
 Focus on what the function does and its main goal.
 2. Arguments: List the parameters, their types, and a brief description
 of what each one represents.
 3. Return: If the function has a return value, describe the return type
 and what it represents. If there's no return, OMIT THE SECTION.
 4. Exceptions: If the function raises any exceptions, list them with descriptions.
 If no exceptions are raised, OMIT THE SECTION.
 5. Side Effects (if applicable): If the function has side effects
 (e.g., modifies global state, interacts with external services), mention them.
 OMIT THE SECTION if it is not lear in the code.
 6. Algorithm or Key Logic (optional): If the function is complex,
 provide a high-level outline of the logic or algorithm involved.
 OMIT THE SECTION if it is not lear in the code.

 Ensure the docstring follows Python's standard docstring conventions and provides
 just enough detail to make the function understandable and usable without overwhelming the reader.
 """

# Get the docstring from OpenAI
docstring=openai_provider.improve_text(prompt, object_code)

# Normalize the docstring by stripping leading/trailing whitespace
normalized_docstring=normalize_docstring(docstring)

# If the docstring is empty, set a fallback docstring
ifnotnormalized_docstring:
print(f"Warning: Generated docstring for {object_type} '{object_name}' is empty. Using a fallback docstring.")
normalized_docstring=f"TODO: Add a docstring for the {object_type} '{object_name}'."

tree=ast.parse(object_code)
indentation_level=0
fornodeinast.walk(tree):
ifisinstance(node, (ast.ClassDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+4# Add 4 spaces for method body
break
ifisinstance(node, (ast.FunctionDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+8# Add 4 spaces for method body
break

# Split the normalized docstring into lines and prefix each with the correct indentation
docstring_lines=normalized_docstring.splitlines()
indented_docstring="\n".join([f"{' '*indentation_level}{line}"forlineindocstring_lines])

# Cache the docstring
cache[entry] = {}
cache[entry]['docstring'] =indented_docstring
cache[entry]['last_hash'] =hashlib.sha256(object_code.encode()).hexdigest()

returnindented_docstring

defvalidate_syntax(code):
"""Validate Python code syntax by writing to a temporary file and checking it."""
try:
# Create a temporary file to check the syntax
withtempfile.NamedTemporaryFile(delete=False, mode='w', suffix=".py") astemp_file:
temp_file.write(code)
temp_filename=temp_file.name

# Run python -m py_compile to check syntax
result=subprocess.run(
 ['python', '-m', 'py_compile', temp_filename],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
 )

ifresult.returncode!=0:
print(code)

os.remove(temp_filename) # Remove the temporary file after checking

ifresult.returncode!=0:
print(f"Syntax error: {result.stderr.decode()}")
returnFalse
returnTrue

exceptExceptionase:
print(f"Error validating syntax: {e}")
returnFalse


defprocess_directory(directory, openai_provider, cache, cache_file):
"""Process all Python files in a directory and its subdirectories."""
forroot, dirs, filesinos.walk(directory):
forfileinfiles:
iffile.endswith(".py"): # Process only Python files
filepath=os.path.join(root, file)
process_file(filepath, openai_provider, cache, cache_file)


defload_cache(cache_file):
"""Load the cache from a file."""
ifos.path.exists(cache_file):
withopen(cache_file, 'r') asf:
returnjson.load(f)
return {}


defsave_cache(cache, cache_file):
"""Save the cache to a file."""
withopen(cache_file, 'w') asf:
json.dump(cache, f, indent=4)

defget_next(tree):
"""Get the next function or class without a docstring."""
# Assuming this function returns the next function/class that lacks a docstring.
# Return None when all functions/classes are documented
objects_without_docstrings=get_functions_and_classes_without_docstrings(tree)
ifobjects_without_docstrings:
returnobjects_without_docstrings[0] # Return the next object
returnNone# No more objects to process

defget_object_node(tree, object_name):
"""
 Traverses the AST to find and return the node for the object (function or class)
 with the specified name.

 Args:
 tree (ast.Module): The AST module to search through.
 object_name (str): The name of the function or class to find.

 Returns:
 ast.FunctionDef | ast.ClassDef: The AST node for the function or class.
 """
fornodeinast.walk(tree):
# Check for function or class definitions by name
ifisinstance(node, (ast.FunctionDef, ast.ClassDef)) andnode.name==object_name:
returnnode
returnNone

defprocess_file(filename, openai_provider, cache, cache_file):
"""Process the file: find functions, classes, generate docstrings, and save the result."""
print(f"Processing file: {filename}")

# Reread the file to ensure we are working with the latest version of the code
withopen(filename, 'r') asfile:
code=file.read()

# Start an infinite loop that will keep processing until there are no more items
whileTrue:
tree=ast.parse(code)
# Get the next function or class that doesn't have a docstring
next_object=get_next(tree)
ifnotnext_object:
print(f"All functions and classes have docstrings in {filename}.")
break# Exit the loop if there are no more functions or classes to document

object_type, object_node=next_object
# Get the object code (e.g., function or class)
object_code=ast.unparse(object_node)
object_name=object_node.name
print(f"Generating docstring for {object_type}: {object_name} in {filename}")

# Check if the function or class already has a docstring
ifast.get_docstring(object_node):
print(f"Skipping {object_type}: {object_name}, it already has a docstring.")
continue

# Generate the docstring using OpenAI or another provider
docstring=generate_docstring(object_code, openai_provider, cache, object_name, object_type, filename)

ifdocstring:
code=insert_docstring_to_object(code, object_node, docstring, object_type)

# Re-read the object code after the docstring insertion
updated_tree=ast.parse(code)
updated_object_node=get_object_node(updated_tree, object_name)
updated_object_code=ast.unparse(updated_object_node)

entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()

# Save the SHA-1 hash of the updated object_code into the cache
cache[entry]['last_hash'] =hashlib.sha1(updated_object_code.encode()).hexdigest()

# Save cache after processing each object
save_cache(cache, cache_file)
print(f"Docstrings added successfully to {object_type}: {object_name} in {filename}.")
else:
print(f"Something wrong with the docstring generation.")

# Validate Python syntax before saving
ifvalidate_syntax(code):
# Write the updated code back to the file
withopen(filename, 'w') asfile: 
file.write(code) 
else:
print(f"Syntax error in {filename}, skipping save.")

deffix_indentation(code):
"""Fix the indentation of the function signature and the code after the docstring."""

# Split the code into lines
code_lines=code.splitlines()

# Find the index where the function signature ends (the first occurrence of ':')
def_index=next(ifori, lineinenumerate(code_lines) ifline.strip().startswith('def '))

# Add 4 spaces to the function signature
code_lines[def_index] =' '+code_lines[def_index].lstrip()

# Now, we need to add indentation for all lines after the docstring
in_docstring=False
foriinrange(def_index+1, len(code_lines)):
line=code_lines[i]

# Detect the start and end of the docstring
ifline.strip().startswith('"""'):
ifin_docstring:
# End of the docstring
in_docstring=False
else:
# Start of the docstring
in_docstring=True

# Indent the line after the docstring
ifin_docstringisFalse:
code_lines[i] =' '+line# Add 4 spaces to the body of the function

# Rejoin the lines into the final code
updated_code="\n".join(code_lines)

returnupdated_code


defremove_duplicate_functions(code):
"""Remove duplicate function definitions based on the function signature."""
# Parse the code into an AST
tree=None
try:
tree=ast.parse(code)
exceptExceptionase:
print(code)
returncode

# Track function signatures
function_signatures=set()
functions_to_remove= []

# Walk through the AST and find function definitions
fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
function_signature=ast.unparse(node)

# If we have seen this function signature before, mark it for removal
iffunction_signatureinfunction_signatures:
functions_to_remove.append(node)
else:
function_signatures.add(function_signature)

# Remove duplicate functions from the code
forfunc_nodeinfunctions_to_remove:
# Removing the function by filtering out the lines corresponding to that function
code=remove_function_from_code(code, func_node)

returncode


defremove_function_from_code(code, function_node):
"""Remove the function from the code."""
# Split the code into lines
code_lines=code.splitlines()

# Get the start and end line numbers of the function
start_line=function_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(ast.unparse(function_node).splitlines())

# Remove the lines of the function from the code
return"\n".join(code_lines[:start_line] +code_lines[end_line+1:])


defextract_main_block(code):
"""Extracts everything from 'if __name__ == "__main__":' to the end of the file."""
lines=code.splitlines() # Split the code into lines
main_start_index=-1# Initialize the index

# Find the line number where '__name__' starts
fori, lineinenumerate(lines):
ifline.strip().startswith('if __name__ == "__main__":'):
main_start_index=i
break

ifmain_start_index==-1:
# print("No '__name__ == \"__main__\"' found in the code.")
return""

# Extract all lines starting from the '__name__' line to the end
return"\n".join(lines[main_start_index:])


definsert_docstring_to_object(code, object_node, docstring, object_type):
main=extract_main_block(code)
"""Insert the generated docstring into the function/class in the correct place."""
object_code=ast.unparse(object_node)

# Get the indentation level of the function or class signature using col_offset
indentation_level=' '* (object_node.col_offset+4) # Add 4 spaces for the body

# Check if the function or class already has a docstring, if so, remove it
existing_docstring=ast.get_docstring(object_node)
ifexisting_docstring:
# Remove the existing docstring from the object_code
object_code=object_code.replace(f'"""{existing_docstring}"""', '')

# The docstring is already properly formatted and indented, so we use it directly
docstring_code=f'{indentation_level}"""\n{docstring}\n{indentation_level}"""'

# Find the first occurrence of ':' to determine where the signature ends
signature_end=object_code.find(':')

func_start_lines=object_code[:signature_end+1]
func_end_lines=object_code[signature_end+1:]

# Insert the docstring immediately after the signature (before the body of the function/class)
object_code_with_docstring=func_start_lines+f'\n{docstring_code}\n'+func_end_lines

ifobject_node.col_offset>0:
object_code_with_docstring=fix_indentation(object_code_with_docstring)

# Get the start and end line numbers of the object_node
start_line=object_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(object_code.splitlines()) -1

# Split the original code into lines
code_lines=code.splitlines()

print("================================")
print(pprint.pprint(code_lines[:start_line]))
print("================================")

# Replace the lines between start_line and end_line with the new code containing the docstring
updated_code_lines= (
code_lines[:start_line] +
 [object_code_with_docstring] +
code_lines[end_line+1:]
 )

# Join the lines back into a single string
updated_code="\n".join(updated_code_lines)
updated_code=remove_duplicate_functions(updated_code)
main_after_update=extract_main_block(updated_code)

ifmain!=main_after_update:
updated_code=updated_code+"\n\n"+main


# Ensure a line break after the inserted docstring for readability
ifnotupdated_code.endswith('\n'):
updated_code+='\n'

returnupdated_code


if__name__=='__main__':
# Get the directory or file from the command line argument
iflen(sys.argv) !=2:
print("Usage: python add_docstrings.py <python_file_or_directory>")
sys.exit(1)

path=sys.argv[1]

# Initialize the OpenAIProvider
openai_provider=OpenAIProvider()

# Load the cache from the cache file (if it exists)
cache_file="docstring_cache.json"
cache=load_cache(cache_file)

ifos.path.isdir(path): # If it's a directory, process it recursively
process_directory(path, openai_provider, cache, cache_file)
elifos.path.isfile(path): # If it's a single file, process it
process_file(path, openai_provider, cache, cache_file)
else:
print(f"The provided path '{path}' is not valid.")
sys.exit(1)

# Save the updated cache to the file
save_cache(cache, cache_file)
#!/usr/bin/python
importast
importos
importrequests
importsys
importjson
importtempfile
importsubprocess
importhashlib
importpprint
importjson


classOpenAIProvider:
def__init__(self):
self.api_key=os.getenv("AI_API_KEY")
ifnotself.api_key:
raiseEnvironmentError("AI_API_KEY not set in environment.")
self.endpoint="https://api.openai.com/v1/chat/completions"
self.model=os.getenv("OPENAI_MODEL", "gpt-4o-mini")

defimprove_text(self, prompt: str, text: str) ->str:
headers= {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
 }

body= {
"model": self.model,
"messages": [
 {"role": "system", "content": prompt},
 {"role": "user", "content": text},
 ],
"temperature": 0.3,
 }

response=requests.post(self.endpoint, json=body, headers=headers, timeout=60)
ifresponse.status_code==200:
returnresponse.json()["choices"][0]["message"]["content"].strip()

raiseException(
f"OpenAI API call failed: {response.status_code} - {response.text}"
 )


defnormalize_docstring(docstring):
"""Normalize docstring to strip leading/trailing spaces, remove markdown, and preserve multi-line format."""
ifnotdocstring:
return""

# Extract everything between triple quotes
# This assumes the docstring is surrounded by triple quotes
start=docstring.find('"""') +3# Skip past the first triple quotes
end=docstring.rfind('"""') # Find the last triple quotes

# Extract the content between the triple quotes
ifstart>2andend>start:
docstring=docstring[start:end].strip() # Remove the triple quotes and extra spaces

returndocstring


defget_functions_and_classes_without_docstrings(tree):
"""Parse the file and return functions and classes without docstrings, 
 considering top-level functions and functions inside a class only."""

functions_and_classes_without_docstrings= []

# We need to track the current class level to identify nested functions
class_stack= []

fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
# Determine the level of indentation
indentation_level=node.col_offset

# Debugging: Print function name, indentation level, and docstring presence
docstring=ast.get_docstring(node)
docstring_status="No docstring"ifdocstringisNoneelse"Has docstring"

# Skip functions with indentation level of 8 (nested functions inside other functions)
ifindentation_level==8:
# print(f"Skipping nested function: {node.name}")
continue
else:
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('function', node))

elifisinstance(node, ast.ClassDef):
# Add the class to the class stack, as we're inside a class now
class_stack.append(node)

# Check if the class has a docstring
docstring=ast.get_docstring(node)
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('class', node))

# Pop the class stack when leaving a class body
ifisinstance(node, ast.ClassDef):
class_stack.pop()

returnfunctions_and_classes_without_docstrings


defgenerate_docstring(object_code, openai_provider, cache, object_name, object_type, filename):
"""Generate a docstring using OpenAIProvider."""
code_digest=hashlib.sha256(object_code.encode()).hexdigest()
entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()
ifentryincacheandcode_digest==cache[entry]['last_hash']:
print(f"Using cached docstring for {object_type}: {object_name}")
returncache[entry]['docstring'] # Return cached docstring if it exists

ifentryincacheandcode_digest!=cache[entry]['last_hash']:
print(f"Has for {object_type}: {object_name} changed. Getting new docstring.")

prompt=f"""
 You will be provided with a Python {object_type} (function or class),
 including its code. Based on this, your task is to generate a Python docstring for it.
 Use the function signature and body to infer the purpose of the function/class,
 the arguments it takes, the return value, and any exceptions it may raise.

 Follow these guidelines to create the docstring:

 1. Summary: Provide a concise summary of the function/class's purpose.
 Focus on what the function does and its main goal.
 2. Arguments: List the parameters, their types, and a brief description
 of what each one represents.
 3. Return: If the function has a return value, describe the return type
 and what it represents. If there's no return, OMIT THE SECTION.
 4. Exceptions: If the function raises any exceptions, list them with descriptions.
 If no exceptions are raised, OMIT THE SECTION.
 5. Side Effects (if applicable): If the function has side effects
 (e.g., modifies global state, interacts with external services), mention them.
 OMIT THE SECTION if it is not lear in the code.
 6. Algorithm or Key Logic (optional): If the function is complex,
 provide a high-level outline of the logic or algorithm involved.
 OMIT THE SECTION if it is not lear in the code.

 Ensure the docstring follows Python's standard docstring conventions and provides
 just enough detail to make the function understandable and usable without overwhelming the reader.
 """

# Get the docstring from OpenAI
docstring=openai_provider.improve_text(prompt, object_code)

# Normalize the docstring by stripping leading/trailing whitespace
normalized_docstring=normalize_docstring(docstring)

# If the docstring is empty, set a fallback docstring
ifnotnormalized_docstring:
print(f"Warning: Generated docstring for {object_type} '{object_name}' is empty. Using a fallback docstring.")
normalized_docstring=f"TODO: Add a docstring for the {object_type} '{object_name}'."

tree=ast.parse(object_code)
indentation_level=0
fornodeinast.walk(tree):
ifisinstance(node, (ast.ClassDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+4# Add 4 spaces for method body
break
ifisinstance(node, (ast.FunctionDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+8# Add 4 spaces for method body
break

# Split the normalized docstring into lines and prefix each with the correct indentation
docstring_lines=normalized_docstring.splitlines()
indented_docstring="\n".join([f"{' '*indentation_level}{line}"forlineindocstring_lines])

# Cache the docstring
cache[entry] = {}
cache[entry]['docstring'] =indented_docstring
cache[entry]['last_hash'] =hashlib.sha256(object_code.encode()).hexdigest()

returnindented_docstring

defvalidate_syntax(code):
"""Validate Python code syntax by writing to a temporary file and checking it."""
try:
# Create a temporary file to check the syntax
withtempfile.NamedTemporaryFile(delete=False, mode='w', suffix=".py") astemp_file:
temp_file.write(code)
temp_filename=temp_file.name

# Run python -m py_compile to check syntax
result=subprocess.run(
 ['python', '-m', 'py_compile', temp_filename],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
 )

ifresult.returncode!=0:
print(code)

os.remove(temp_filename) # Remove the temporary file after checking

ifresult.returncode!=0:
print(f"Syntax error: {result.stderr.decode()}")
returnFalse
returnTrue

exceptExceptionase:
print(f"Error validating syntax: {e}")
returnFalse


defprocess_directory(directory, openai_provider, cache, cache_file):
"""Process all Python files in a directory and its subdirectories."""
forroot, dirs, filesinos.walk(directory):
forfileinfiles:
iffile.endswith(".py"): # Process only Python files
filepath=os.path.join(root, file)
process_file(filepath, openai_provider, cache, cache_file)


defload_cache(cache_file):
"""Load the cache from a file."""
ifos.path.exists(cache_file):
withopen(cache_file, 'r') asf:
returnjson.load(f)
return {}


defsave_cache(cache, cache_file):
"""Save the cache to a file."""
withopen(cache_file, 'w') asf:
json.dump(cache, f, indent=4)

defget_next(tree):
"""Get the next function or class without a docstring."""
# Assuming this function returns the next function/class that lacks a docstring.
# Return None when all functions/classes are documented
objects_without_docstrings=get_functions_and_classes_without_docstrings(tree)
ifobjects_without_docstrings:
returnobjects_without_docstrings[0] # Return the next object
returnNone# No more objects to process

defget_object_node(tree, object_name):
"""
 Traverses the AST to find and return the node for the object (function or class)
 with the specified name.

 Args:
 tree (ast.Module): The AST module to search through.
 object_name (str): The name of the function or class to find.

 Returns:
 ast.FunctionDef | ast.ClassDef: The AST node for the function or class.
 """
fornodeinast.walk(tree):
# Check for function or class definitions by name
ifisinstance(node, (ast.FunctionDef, ast.ClassDef)) andnode.name==object_name:
returnnode
returnNone

defprocess_file(filename, openai_provider, cache, cache_file):
"""Process the file: find functions, classes, generate docstrings, and save the result."""
print(f"Processing file: {filename}")

# Reread the file to ensure we are working with the latest version of the code
withopen(filename, 'r') asfile:
code=file.read()

# Start an infinite loop that will keep processing until there are no more items
whileTrue:
tree=ast.parse(code)
# Get the next function or class that doesn't have a docstring
next_object=get_next(tree)
ifnotnext_object:
print(f"All functions and classes have docstrings in {filename}.")
break# Exit the loop if there are no more functions or classes to document

object_type, object_node=next_object
# Get the object code (e.g., function or class)
object_code=ast.unparse(object_node)
object_name=object_node.name
print(f"Generating docstring for {object_type}: {object_name} in {filename}")

# Check if the function or class already has a docstring
ifast.get_docstring(object_node):
print(f"Skipping {object_type}: {object_name}, it already has a docstring.")
continue

# Generate the docstring using OpenAI or another provider
docstring=generate_docstring(object_code, openai_provider, cache, object_name, object_type, filename)

ifdocstring:
code=insert_docstring_to_object(code, object_node, docstring, object_type)

# Re-read the object code after the docstring insertion
updated_tree=ast.parse(code)
updated_object_node=get_object_node(updated_tree, object_name)
updated_object_code=ast.unparse(updated_object_node)

entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()

# Save the SHA-1 hash of the updated object_code into the cache
cache[entry]['last_hash'] =hashlib.sha1(updated_object_code.encode()).hexdigest()

# Save cache after processing each object
save_cache(cache, cache_file)
print(f"Docstrings added successfully to {object_type}: {object_name} in {filename}.")
else:
print(f"Something wrong with the docstring generation.")

# Validate Python syntax before saving
ifvalidate_syntax(code):
# Write the updated code back to the file
withopen(filename, 'w') asfile: 
file.write(code) 
else:
print(f"Syntax error in {filename}, skipping save.")

deffix_indentation(code):
"""Fix the indentation of the function signature and the code after the docstring."""

# Split the code into lines
code_lines=code.splitlines()

# Find the index where the function signature ends (the first occurrence of ':')
def_index=next(ifori, lineinenumerate(code_lines) ifline.strip().startswith('def '))

# Add 4 spaces to the function signature
code_lines[def_index] =' '+code_lines[def_index].lstrip()

# Now, we need to add indentation for all lines after the docstring
in_docstring=False
foriinrange(def_index+1, len(code_lines)):
line=code_lines[i]

# Detect the start and end of the docstring
ifline.strip().startswith('"""'):
ifin_docstring:
# End of the docstring
in_docstring=False
else:
# Start of the docstring
in_docstring=True

# Indent the line after the docstring
ifin_docstringisFalse:
code_lines[i] =' '+line# Add 4 spaces to the body of the function

# Rejoin the lines into the final code
updated_code="\n".join(code_lines)

returnupdated_code


defremove_duplicate_functions(code):
"""Remove duplicate function definitions based on the function signature."""
# Parse the code into an AST
tree=None
try:
tree=ast.parse(code)
exceptExceptionase:
print(code)
returncode

# Track function signatures
function_signatures=set()
functions_to_remove= []

# Walk through the AST and find function definitions
fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
function_signature=ast.unparse(node)

# If we have seen this function signature before, mark it for removal
iffunction_signatureinfunction_signatures:
functions_to_remove.append(node)
else:
function_signatures.add(function_signature)

# Remove duplicate functions from the code
forfunc_nodeinfunctions_to_remove:
# Removing the function by filtering out the lines corresponding to that function
code=remove_function_from_code(code, func_node)

returncode


defremove_function_from_code(code, function_node):
"""Remove the function from the code."""
# Split the code into lines
code_lines=code.splitlines()

# Get the start and end line numbers of the function
start_line=function_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(ast.unparse(function_node).splitlines())

# Remove the lines of the function from the code
return"\n".join(code_lines[:start_line] +code_lines[end_line+1:])


defextract_main_block(code):
"""Extracts everything from 'if __name__ == "__main__":' to the end of the file."""
lines=code.splitlines() # Split the code into lines
main_start_index=-1# Initialize the index

# Find the line number where '__name__' starts
fori, lineinenumerate(lines):
ifline.strip().startswith('if __name__ == "__main__":'):
main_start_index=i
break

ifmain_start_index==-1:
# print("No '__name__ == \"__main__\"' found in the code.")
return""

# Extract all lines starting from the '__name__' line to the end
return"\n".join(lines[main_start_index:])


definsert_docstring_to_object(code, object_node, docstring, object_type):
main=extract_main_block(code)
"""Insert the generated docstring into the function/class in the correct place."""
object_code=ast.unparse(object_node)

# Get the indentation level of the function or class signature using col_offset
indentation_level=' '* (object_node.col_offset+4) # Add 4 spaces for the body

# Check if the function or class already has a docstring, if so, remove it
existing_docstring=ast.get_docstring(object_node)
ifexisting_docstring:
# Remove the existing docstring from the object_code
object_code=object_code.replace(f'"""{existing_docstring}"""', '')

# The docstring is already properly formatted and indented, so we use it directly
docstring_code=f'{indentation_level}"""\n{docstring}\n{indentation_level}"""'

# Find the first occurrence of ':' to determine where the signature ends
signature_end=object_code.find(':')

func_start_lines=object_code[:signature_end+1]
func_end_lines=object_code[signature_end+1:]

# Insert the docstring immediately after the signature (before the body of the function/class)
object_code_with_docstring=func_start_lines+f'\n{docstring_code}\n'+func_end_lines

ifobject_node.col_offset>0:
object_code_with_docstring=fix_indentation(object_code_with_docstring)

# Get the start and end line numbers of the object_node
start_line=object_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(object_code.splitlines()) -1

# Split the original code into lines
code_lines=code.splitlines()

print("================================")
print(pprint.pprint(code_lines[:start_line]))
print("================================")

# Replace the lines between start_line and end_line with the new code containing the docstring
updated_code_lines= (
code_lines[:start_line] +
 [object_code_with_docstring] +
code_lines[end_line+1:]
 )

# Join the lines back into a single string
updated_code="\n".join(updated_code_lines)
updated_code=remove_duplicate_functions(updated_code)
main_after_update=extract_main_block(updated_code)

ifmain!=main_after_update:
updated_code=updated_code+"\n\n"+main


# Ensure a line break after the inserted docstring for readability
ifnotupdated_code.endswith('\n'):
updated_code+='\n'

returnupdated_code


if__name__=='__main__':
# Get the directory or file from the command line argument
iflen(sys.argv) !=2:
print("Usage: python add_docstrings.py <python_file_or_directory>")
sys.exit(1)

path=sys.argv[1]

# Initialize the OpenAIProvider
openai_provider=OpenAIProvider()

# Load the cache from the cache file (if it exists)
cache_file="docstring_cache.json"
cache=load_cache(cache_file)

ifos.path.isdir(path): # If it's a directory, process it recursively
process_directory(path, openai_provider, cache, cache_file)
elifos.path.isfile(path): # If it's a single file, process it
process_file(path, openai_provider, cache, cache_file)
else:
print(f"The provided path '{path}' is not valid.")
sys.exit(1)

# Save the updated cache to the file
save_cache(cache, cache_file)
#!/usr/bin/python
importast
importos
importrequests
importsys
importjson
importtempfile
importsubprocess
importhashlib
importpprint
importjson


classOpenAIProvider:
def__init__(self):
self.api_key=os.getenv("AI_API_KEY")
ifnotself.api_key:
raiseEnvironmentError("AI_API_KEY not set in environment.")
self.endpoint="https://api.openai.com/v1/chat/completions"
self.model=os.getenv("OPENAI_MODEL", "gpt-4o-mini")

defimprove_text(self, prompt: str, text: str) ->str:
headers= {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
 }

body= {
"model": self.model,
"messages": [
 {"role": "system", "content": prompt},
 {"role": "user", "content": text},
 ],
"temperature": 0.3,
 }

response=requests.post(self.endpoint, json=body, headers=headers, timeout=60)
ifresponse.status_code==200:
returnresponse.json()["choices"][0]["message"]["content"].strip()

raiseException(
f"OpenAI API call failed: {response.status_code} - {response.text}"
 )


defnormalize_docstring(docstring):
"""Normalize docstring to strip leading/trailing spaces, remove markdown, and preserve multi-line format."""
ifnotdocstring:
return""

# Extract everything between triple quotes
# This assumes the docstring is surrounded by triple quotes
start=docstring.find('"""') +3# Skip past the first triple quotes
end=docstring.rfind('"""') # Find the last triple quotes

# Extract the content between the triple quotes
ifstart>2andend>start:
docstring=docstring[start:end].strip() # Remove the triple quotes and extra spaces

returndocstring


defget_functions_and_classes_without_docstrings(tree):
"""Parse the file and return functions and classes without docstrings, 
 considering top-level functions and functions inside a class only."""

functions_and_classes_without_docstrings= []

# We need to track the current class level to identify nested functions
class_stack= []

fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
# Determine the level of indentation
indentation_level=node.col_offset

# Debugging: Print function name, indentation level, and docstring presence
docstring=ast.get_docstring(node)
docstring_status="No docstring"ifdocstringisNoneelse"Has docstring"

# Skip functions with indentation level of 8 (nested functions inside other functions)
ifindentation_level==8:
# print(f"Skipping nested function: {node.name}")
continue
else:
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('function', node))

elifisinstance(node, ast.ClassDef):
# Add the class to the class stack, as we're inside a class now
class_stack.append(node)

# Check if the class has a docstring
docstring=ast.get_docstring(node)
ifdocstringisNone:
functions_and_classes_without_docstrings.append(('class', node))

# Pop the class stack when leaving a class body
ifisinstance(node, ast.ClassDef):
class_stack.pop()

returnfunctions_and_classes_without_docstrings


defgenerate_docstring(object_code, openai_provider, cache, object_name, object_type, filename):
"""Generate a docstring using OpenAIProvider."""
code_digest=hashlib.sha256(object_code.encode()).hexdigest()
entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()
ifentryincacheandcode_digest==cache[entry]['last_hash']:
print(f"Using cached docstring for {object_type}: {object_name}")
returncache[entry]['docstring'] # Return cached docstring if it exists

ifentryincacheandcode_digest!=cache[entry]['last_hash']:
print(f"Has for {object_type}: {object_name} changed. Getting new docstring.")

prompt=f"""
 You will be provided with a Python {object_type} (function or class),
 including its code. Based on this, your task is to generate a Python docstring for it.
 Use the function signature and body to infer the purpose of the function/class,
 the arguments it takes, the return value, and any exceptions it may raise.

 Follow these guidelines to create the docstring:

 1. Summary: Provide a concise summary of the function/class's purpose.
 Focus on what the function does and its main goal.
 2. Arguments: List the parameters, their types, and a brief description
 of what each one represents.
 3. Return: If the function has a return value, describe the return type
 and what it represents. If there's no return, OMIT THE SECTION.
 4. Exceptions: If the function raises any exceptions, list them with descriptions.
 If no exceptions are raised, OMIT THE SECTION.
 5. Side Effects (if applicable): If the function has side effects
 (e.g., modifies global state, interacts with external services), mention them.
 OMIT THE SECTION if it is not lear in the code.
 6. Algorithm or Key Logic (optional): If the function is complex,
 provide a high-level outline of the logic or algorithm involved.
 OMIT THE SECTION if it is not lear in the code.

 Ensure the docstring follows Python's standard docstring conventions and provides
 just enough detail to make the function understandable and usable without overwhelming the reader.
 """

# Get the docstring from OpenAI
docstring=openai_provider.improve_text(prompt, object_code)

# Normalize the docstring by stripping leading/trailing whitespace
normalized_docstring=normalize_docstring(docstring)

# If the docstring is empty, set a fallback docstring
ifnotnormalized_docstring:
print(f"Warning: Generated docstring for {object_type} '{object_name}' is empty. Using a fallback docstring.")
normalized_docstring=f"TODO: Add a docstring for the {object_type} '{object_name}'."

tree=ast.parse(object_code)
indentation_level=0
fornodeinast.walk(tree):
ifisinstance(node, (ast.ClassDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+4# Add 4 spaces for method body
break
ifisinstance(node, (ast.FunctionDef)):
ifnode.name==object_name:
indentation_level=node.col_offset+8# Add 4 spaces for method body
break

# Split the normalized docstring into lines and prefix each with the correct indentation
docstring_lines=normalized_docstring.splitlines()
indented_docstring="\n".join([f"{' '*indentation_level}{line}"forlineindocstring_lines])

# Cache the docstring
cache[entry] = {}
cache[entry]['docstring'] =indented_docstring
cache[entry]['last_hash'] =hashlib.sha256(object_code.encode()).hexdigest()

returnindented_docstring

defvalidate_syntax(code):
"""Validate Python code syntax by writing to a temporary file and checking it."""
try:
# Create a temporary file to check the syntax
withtempfile.NamedTemporaryFile(delete=False, mode='w', suffix=".py") astemp_file:
temp_file.write(code)
temp_filename=temp_file.name

# Run python -m py_compile to check syntax
result=subprocess.run(
 ['python', '-m', 'py_compile', temp_filename],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
 )

ifresult.returncode!=0:
print(code)

os.remove(temp_filename) # Remove the temporary file after checking

ifresult.returncode!=0:
print(f"Syntax error: {result.stderr.decode()}")
returnFalse
returnTrue

exceptExceptionase:
print(f"Error validating syntax: {e}")
returnFalse


defprocess_directory(directory, openai_provider, cache, cache_file):
"""Process all Python files in a directory and its subdirectories."""
forroot, dirs, filesinos.walk(directory):
forfileinfiles:
iffile.endswith(".py"): # Process only Python files
filepath=os.path.join(root, file)
process_file(filepath, openai_provider, cache, cache_file)


defload_cache(cache_file):
"""Load the cache from a file."""
ifos.path.exists(cache_file):
withopen(cache_file, 'r') asf:
returnjson.load(f)
return {}


defsave_cache(cache, cache_file):
"""Save the cache to a file."""
withopen(cache_file, 'w') asf:
json.dump(cache, f, indent=4)

defget_next(tree):
"""Get the next function or class without a docstring."""
# Assuming this function returns the next function/class that lacks a docstring.
# Return None when all functions/classes are documented
objects_without_docstrings=get_functions_and_classes_without_docstrings(tree)
ifobjects_without_docstrings:
returnobjects_without_docstrings[0] # Return the next object
returnNone# No more objects to process

defget_object_node(tree, object_name):
"""
 Traverses the AST to find and return the node for the object (function or class)
 with the specified name.

 Args:
 tree (ast.Module): The AST module to search through.
 object_name (str): The name of the function or class to find.

 Returns:
 ast.FunctionDef | ast.ClassDef: The AST node for the function or class.
 """
fornodeinast.walk(tree):
# Check for function or class definitions by name
ifisinstance(node, (ast.FunctionDef, ast.ClassDef)) andnode.name==object_name:
returnnode
returnNone

defprocess_file(filename, openai_provider, cache, cache_file):
"""Process the file: find functions, classes, generate docstrings, and save the result."""
print(f"Processing file: {filename}")

# Reread the file to ensure we are working with the latest version of the code
withopen(filename, 'r') asfile:
code=file.read()

# Start an infinite loop that will keep processing until there are no more items
whileTrue:
tree=ast.parse(code)
# Get the next function or class that doesn't have a docstring
next_object=get_next(tree)
ifnotnext_object:
print(f"All functions and classes have docstrings in {filename}.")
break# Exit the loop if there are no more functions or classes to document

object_type, object_node=next_object
# Get the object code (e.g., function or class)
object_code=ast.unparse(object_node)
object_name=object_node.name
print(f"Generating docstring for {object_type}: {object_name} in {filename}")

# Check if the function or class already has a docstring
ifast.get_docstring(object_node):
print(f"Skipping {object_type}: {object_name}, it already has a docstring.")
continue

# Generate the docstring using OpenAI or another provider
docstring=generate_docstring(object_code, openai_provider, cache, object_name, object_type, filename)

ifdocstring:
code=insert_docstring_to_object(code, object_node, docstring, object_type)

# Re-read the object code after the docstring insertion
updated_tree=ast.parse(code)
updated_object_node=get_object_node(updated_tree, object_name)
updated_object_code=ast.unparse(updated_object_node)

entry=filename+"-"+object_name
entry=hashlib.sha256(entry.encode()).hexdigest()

# Save the SHA-1 hash of the updated object_code into the cache
cache[entry]['last_hash'] =hashlib.sha1(updated_object_code.encode()).hexdigest()

# Save cache after processing each object
save_cache(cache, cache_file)
print(f"Docstrings added successfully to {object_type}: {object_name} in {filename}.")
else:
print(f"Something wrong with the docstring generation.")

# Validate Python syntax before saving
ifvalidate_syntax(code):
# Write the updated code back to the file
withopen(filename, 'w') asfile: 
file.write(code) 
else:
print(f"Syntax error in {filename}, skipping save.")

deffix_indentation(code):
"""Fix the indentation of the function signature and the code after the docstring."""

# Split the code into lines
code_lines=code.splitlines()

# Find the index where the function signature ends (the first occurrence of ':')
def_index=next(ifori, lineinenumerate(code_lines) ifline.strip().startswith('def '))

# Add 4 spaces to the function signature
code_lines[def_index] =' '+code_lines[def_index].lstrip()

# Now, we need to add indentation for all lines after the docstring
in_docstring=False
foriinrange(def_index+1, len(code_lines)):
line=code_lines[i]

# Detect the start and end of the docstring
ifline.strip().startswith('"""'):
ifin_docstring:
# End of the docstring
in_docstring=False
else:
# Start of the docstring
in_docstring=True

# Indent the line after the docstring
ifin_docstringisFalse:
code_lines[i] =' '+line# Add 4 spaces to the body of the function

# Rejoin the lines into the final code
updated_code="\n".join(code_lines)

returnupdated_code


defremove_duplicate_functions(code):
"""Remove duplicate function definitions based on the function signature."""
# Parse the code into an AST
tree=None
try:
tree=ast.parse(code)
exceptExceptionase:
print(code)
returncode

# Track function signatures
function_signatures=set()
functions_to_remove= []

# Walk through the AST and find function definitions
fornodeinast.walk(tree):
ifisinstance(node, ast.FunctionDef):
function_signature=ast.unparse(node)

# If we have seen this function signature before, mark it for removal
iffunction_signatureinfunction_signatures:
functions_to_remove.append(node)
else:
function_signatures.add(function_signature)

# Remove duplicate functions from the code
forfunc_nodeinfunctions_to_remove:
# Removing the function by filtering out the lines corresponding to that function
code=remove_function_from_code(code, func_node)

returncode


defremove_function_from_code(code, function_node):
"""Remove the function from the code."""
# Split the code into lines
code_lines=code.splitlines()

# Get the start and end line numbers of the function
start_line=function_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(ast.unparse(function_node).splitlines())

# Remove the lines of the function from the code
return"\n".join(code_lines[:start_line] +code_lines[end_line+1:])


defextract_main_block(code):
"""Extracts everything from 'if __name__ == "__main__":' to the end of the file."""
lines=code.splitlines() # Split the code into lines
main_start_index=-1# Initialize the index

# Find the line number where '__name__' starts
fori, lineinenumerate(lines):
ifline.strip().startswith('if __name__ == "__main__":'):
main_start_index=i
break

ifmain_start_index==-1:
# print("No '__name__ == \"__main__\"' found in the code.")
return""

# Extract all lines starting from the '__name__' line to the end
return"\n".join(lines[main_start_index:])


definsert_docstring_to_object(code, object_node, docstring, object_type):
main=extract_main_block(code)
"""Insert the generated docstring into the function/class in the correct place."""
object_code=ast.unparse(object_node)

# Get the indentation level of the function or class signature using col_offset
indentation_level=' '* (object_node.col_offset+4) # Add 4 spaces for the body

# Check if the function or class already has a docstring, if so, remove it
existing_docstring=ast.get_docstring(object_node)
ifexisting_docstring:
# Remove the existing docstring from the object_code
object_code=object_code.replace(f'"""{existing_docstring}"""', '')

# The docstring is already properly formatted and indented, so we use it directly
docstring_code=f'{indentation_level}"""\n{docstring}\n{indentation_level}"""'

# Find the first occurrence of ':' to determine where the signature ends
signature_end=object_code.find(':')

func_start_lines=object_code[:signature_end+1]
func_end_lines=object_code[signature_end+1:]

# Insert the docstring immediately after the signature (before the body of the function/class)
object_code_with_docstring=func_start_lines+f'\n{docstring_code}\n'+func_end_lines

ifobject_node.col_offset>0:
object_code_with_docstring=fix_indentation(object_code_with_docstring)

# Get the start and end line numbers of the object_node
start_line=object_node.lineno-1# Adjust for 0-indexing
end_line=start_line+len(object_code.splitlines()) -1

# Split the original code into lines
code_lines=code.splitlines()

print("================================")
print(pprint.pprint(code_lines[:start_line]))
print("================================")

# Replace the lines between start_line and end_line with the new code containing the docstring
updated_code_lines= (
code_lines[:start_line] +
 [object_code_with_docstring] +
code_lines[end_line+1:]
 )

# Join the lines back into a single string
updated_code="\n".join(updated_code_lines)
updated_code=remove_duplicate_functions(updated_code)
main_after_update=extract_main_block(updated_code)

ifmain!=main_after_update:
updated_code=updated_code+"\n\n"+main


# Ensure a line break after the inserted docstring for readability
ifnotupdated_code.endswith('\n'):
updated_code+='\n'

returnupdated_code


if__name__=='__main__':
# Get the directory or file from the command line argument
iflen(sys.argv) !=2:
print("Usage: python add_docstrings.py <python_file_or_directory>")
sys.exit(1)

path=sys.argv[1]

# Initialize the OpenAIProvider
openai_provider=OpenAIProvider()

# Load the cache from the cache file (if it exists)
cache_file="docstring_cache.json"
cache=load_cache(cache_file)

ifos.path.isdir(path): # If it's a directory, process it recursively
process_directory(path, openai_provider, cache, cache_file)
elifos.path.isfile(path): # If it's a single file, process it
process_file(path, openai_provider, cache, cache_file)
else:
print(f"The provided path '{path}' is not valid.")
sys.exit(1)

# Save the updated cache to the file
save_cache(cache, cache_file)