A module for generating AI-based code documentation and data flow diagrams.
Project description
FastWrite
Python Module for AI-Assisted Documentation
Current Statistics:
Overview
This module provides functionality to:
- Process Code Files: Extract and list Python files from a ZIP archive.
- Generate Data Flow Diagrams: Create a data flow chart (in Graphviz format) by analyzing Python code using the AST module.
- Generate Documentation: Produce detailed documentation for Python code using multiple AI models:
- Groq-based models (remote)
- Gemini-based models (remote)
- OpenAI-based models (remote)
- OpenRouter-based models (remote)
- Ollama-based models (local)
- Evaluate Documentation Quality: Compute BLEU/ROUGE scores to compare generated documentation against a reference document.
Installation
Requirements
- Python 3.11
- groq
- google-genai
- openai
- requests
- nltk
- python-dotenv
- rouge_metric
Install Dependencies
pip install groq google-genai requests nltk python-dotenv openai rouge_metric
Usage
Processing Files:
from FastWrite import extract_zip, list_python_files, read_file
import tempfile
import os
# Specify the path to your ZIP file containing Python code
zip_file_path = "path/to/your/code.zip"
with tempfile.TemporaryDirectory() as tmp_dir:
# Extract the ZIP file
extract_zip(zip_file_path, tmp_dir)
# List Python files in the extracted directory
py_files = list_python_files(tmp_dir)
if py_files:
# For example, choose the first Python file as the main file
main_file_path = os.path.join(tmp_dir, py_files[0])
code_content = read_file(main_file_path)
Generating Data Flow Diagrams:
from FastWrite import generate_data_flow
# Generate Graphviz code for the data flow diagram
graphviz_code = generate_data_flow(code_content)
print(graphviz_code)
Generating Documentation (Express Mode):
# Normal Mode:
py -m FastWrite --LLM_NAME
# With custom code file
py -m FastWrite code_filename.py --LLM_NAME
# Simplified (Novice Friendly) Mode:
py -m FastWrite code_filename.py --LLM_NAME --Simplify
# Formal (Precise and to the Point) Mode:
py -m FastWrite code_filename.py --LLM_NAME --Formal
# Research (In-Depth and Elaborate Documentation):
py -m FastWrite code_filename.py --LLM_NAME --Research
# Custom (Allows you to feed a custom prompt of your own):
py -m FastWrite code_filename.py --LLM_NAME --Custom-Prompt "Enter Custom Prompt here"
Generating Documentation (Groq):
from FastWrite import generate_documentation_groq
custom_prompt = """
Objective:
Generate high-quality, developer-friendly documentation for the following Python code Ensure you include Detailed function-level and file-level documentation and a high level slightly less technical documentation at the start to make it friendly. Do not print full code snippets of existing code, just explain them:
"""
groq_api_key = "your_groq_api_key"
groq_model = "deepseek-r1-distill-llama-70b" # Replace with your desired model
doc_groq = generate_documentation_groq(code_content, custom_prompt, groq_api_key, groq_model)
print(doc_groq)
Generating Documentation (Gemini):
from FastWrite import generate_documentation_gemini
custom_prompt = """
Objective:
Generate high-quality, developer-friendly documentation for the following Python code Ensure you include Detailed function-level and file-level documentation and a high level slightly less technical documentation at the start to make it friendly. Do not print full code snippets of existing code, just explain them:
"""
gemini_api_key = "your_gemini_api_key"
gemini_model = "gemini-2.0-flash" # Replace with your desired model
doc_gemini = generate_documentation_gemini(code_content, custom_prompt, gemini_api_key, gemini_model)
print(doc_gemini)
Generating Documentation (OpenAI):
from FastWrite import generate_documentation_openai
custom_prompt = """
Objective:
Generate high-quality, developer-friendly documentation for the following Python code Ensure you include Detailed function-level and file-level documentation and a high level slightly less technical documentation at the start to make it friendly. Do not print full code snippets of existing code, just explain them:
"""
doc_openai = generate_documentation_openai(code_content, custom_prompt)
print(doc_openai)
Generating Documentation (Ollama):
from FastWrite import generate_documentation_ollama
custom_prompt = """
Objective:
Generate high-quality, developer-friendly documentation for the following Python code Ensure you include Detailed function-level and file-level documentation and a high level slightly less technical documentation at the start to make it friendly. Do not print full code snippets of existing code, just explain them:
"""
# Replace with your local Ollama model name (e.g., "ollama-llama-70b")
ollama_model = "ollama-llama-70b"
doc_ollama = generate_documentation_ollama(code_content, custom_prompt, ollama_model)
print(doc_ollama)
Generating Documentation (OpenRouter):
from FastWrite import generate_documentation_openrouter
custom_prompt = """
Objective:
Generate high-quality, developer-friendly documentation for the following Python code Ensure you include Detailed function-level and file-level documentation and a high level slightly less technical documentation at the start to make it friendly. Do not print full code snippets of existing code, just explain them:
"""
doc_openrouter = generate_documentation_openrouter(code_content, custom_prompt)
print(doc_openrouter)
Calculating Bleu Score:
from FastWrite import calculate_bleu
# Provide a reference documentation string for comparison
reference_doc = "Your reference documentation text here..."
# Basic BLEU score calculation
bleu_score = calculate_bleu(doc_llm_host, reference_doc) ##LLM host may include Groq,Gemini,OpenAI or Ollama
print("BLEU Score:", bleu_score)
# BLEU score with smoothing
# Available smoothing methods: 'method0', 'method1', 'method2', 'method3', 'method4', 'method5', 'method6', 'method7'
bleu_score_smoothed = calculate_bleu(doc_llm_host, reference_doc, smoothing_method='method1')
print("BLEU Score with Smoothing:", bleu_score_smoothed)
# Calculate BLEU score with multiple reference documents
from FastWrite import calculate_bleu_multi_reference
reference_docs = ["Reference doc 1...", "Reference doc 2..."]
bleu_score_multi = calculate_bleu_multi_reference(doc_llm_host, reference_docs, smoothing_method='method1')
print("BLEU Score with Multiple References:", bleu_score_multi)
# Calculate BLEU score directly from files
from FastWrite import calculate_bleu_from_files
bleu_score_from_files = calculate_bleu_from_files("candidate_doc.txt", "reference_doc.txt", smoothing_method='method1')
print("BLEU Score from Files:", bleu_score_from_files)
Calculating Rouge Score:
from FastWrite.rouge import calculate_rouge
# Provide a reference documentation string for comparison
reference_doc = "Your reference documentation text here..."
# Calculate Rouge
rouge = calculate_rouge(doc_llm_host,reference_doc)
Generating README File:
from FastWrite.print import readmegen
readmegen(doc_llm,llm_used)
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
fastwrite-1.2.1.tar.gz
(11.1 kB
view details)
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
fastwrite-1.2.1-py3-none-any.whl
(13.9 kB
view details)
File details
Details for the file fastwrite-1.2.1.tar.gz.
File metadata
- Download URL: fastwrite-1.2.1.tar.gz
- Upload date:
- Size: 11.1 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.1.0 CPython/3.13.3
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
77dfcae96904afc921fba77f4dab80671981ff0a18b98324ce127c6119e48f7f
|
|
| MD5 |
a5a5b7f86ce26706fc3058446a2c7998
|
|
| BLAKE2b-256 |
5abdc3639d5efa7590eac61215c655b14d26e8a10d8b2a05e228f291ccd8386f
|
File details
Details for the file fastwrite-1.2.1-py3-none-any.whl.
File metadata
- Download URL: fastwrite-1.2.1-py3-none-any.whl
- Upload date:
- Size: 13.9 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.1.0 CPython/3.13.3
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
772ca2e820b8235e94352269a4f97cbb848e8c081d426b36fe630c99b52d2bc9
|
|
| MD5 |
d8c6bb31d3e92cd34d446c7cf6af1c97
|
|
| BLAKE2b-256 |
f85a58bc50fdd6496c4cc7e97e26518f0bb1bd2227f28d00f7e0fc5f7c47e019
|