doc_generator.py

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")