File size: 4,497 Bytes
d12bff5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
import google.generativeai as genai
import re
import os
import ast
from dotenv import load_dotenv
import sys
import importlib.util
load_dotenv()
API_KEY = os.getenv("GOOGLE_API_KEY")
if API_KEY is None:
raise ValueError("⚠️ The API key MY_API_KEY is missing! Check the Secrets in Hugging Face.")
genai.configure(api_key=API_KEY)
model = genai.GenerativeModel("models/gemini-2.0-flash")
PROMPT = """You are an expert programming assistant.
For the following code, perform the following actions:
- The code must remain exactly the same
- Add clear comments for each important step.
- Rename variables if it makes the code easier to understand.
- Add type annotations if the language supports it.
- For each function, add a Google-style docstring (or equivalent format depending on the language).
Respond only with the updated code, no explanation.
Here is the code:
{code}
"""
def generate_documented_code(input_path: str, output_path: str) -> str:
"""
Generate a documented version of the code from the given input file and save it to the output file.
Args:
input_path (str): Path to the original code file.
output_path (str): Path where the documented code will be saved.
Returns:
str: The updated and documented code.
"""
with open(input_path, "r", encoding="utf-8") as f:
original_code = f.read()
prompt = PROMPT.format(code=original_code)
response = model.generate_content(prompt)
updated_code = response.text.strip()
# Clean up Markdown blocks if present
lines = updated_code.splitlines()
if len(lines) > 2:
lines = lines[1:-1] # remove the first and last lines
updated_code = "\n".join(lines)
else:
# if less than 3 lines, clear everything or keep as is depending on needs
updated_code = ""
with open(output_path, "w", encoding="utf-8") as output_file:
output_file.write(updated_code)
return updated_code
def extract_imports_from_file(file_path):
"""
Extract imported modules from a Python file to generate requirements.txt.
Args:
file_path (str): Path to the Python file.
Returns:
set: A set of imported module names.
"""
try:
with open(file_path, "r", encoding="utf-8") as f:
tree = ast.parse(f.read())
except SyntaxError:
return set()
imports = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
imports.add(alias.name.split('.')[0])
elif isinstance(node, ast.ImportFrom):
if node.module and not node.module.startswith("."):
imports.add(node.module.split('.')[0])
return imports
def is_std_lib(module_name):
"""
Check if a module is part of the Python standard library.
Args:
module_name (str): The name of the module.
Returns:
bool: True if the module is part of the standard library, False otherwise.
"""
if module_name in sys.builtin_module_names:
return True
spec = importlib.util.find_spec(module_name)
return spec is not None and "site-packages" not in (spec.origin or "")
def generate_requirements_txt(base_path, output_path):
"""
Generate a requirements.txt file based on external imports found in Python files.
Args:
base_path (str): Root directory of the codebase.
output_path (str): Path to save the generated requirements.txt file.
"""
all_imports = set()
local_modules = set()
# Get names of internal modules (i.e., .py files in the repo)
for root, _, files in os.walk(base_path):
for file in files:
if file.endswith(".py"):
module_name = os.path.splitext(file)[0]
local_modules.add(module_name)
# Extract all imports used in the project
for root, _, files in os.walk(base_path):
for file in files:
if file.endswith(".py"):
file_path = os.path.join(root, file)
all_imports.update(extract_imports_from_file(file_path))
# Remove internal modules and standard library modules
external_imports = sorted([
imp for imp in all_imports
if imp not in local_modules and not is_std_lib(imp)
])
# Write the requirements.txt file
with open(output_path, "w", encoding="utf-8") as f:
for package in external_imports:
f.write(f"{package}\n")
|