#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PythonコードにSphinx形式のDocstringを自動追加するツール。
このスクリプトは、指定されたPythonソースファイルの内容をAIモデルに送信し、
Sphinx形式のDocstringが挿入された完成版のコードを生成します。
コマンドライン引数で入力ファイルパターン、出力ファイル名、
INIファイルによるプロンプト設定、使用するAIモデルなどを指定できます。
生成されたDocstringは、関数やクラスの概要、詳細説明、引数、戻り値などを
日本語で記述します。既存のファイルを上書きするか、更新日時でスキップするかなど、
柔軟なファイル処理オプションを提供します。
:doc:`add_docstring_usage`
"""
import os
import sys
import argparse
import glob
import time
import re
import traceback
from pathlib import Path
# AI連携用ライブラリのインポート
try:
from tkai_lib import read_ai_config
from tkai_lib import query_openai4, query_openai5, query_google
from tkai_lib import extract_openai5_text
except ImportError:
print("Error: tkai_lib.py が見つかりません。ライブラリのパスを確認してください。", file=sys.stderr)
sys.exit(1)
# =========================================================
# INI検索・読み込みロジック (explain_program5.py の設計を継承)
# =========================================================
[ドキュメント]
def search_file(infile=None):
"""
指定されたファイルを作業ディレクトリとスクリプトディレクトリから検索する。
`infile` が指定されない場合、実行スクリプト名からデフォルトのINIファイル名を生成します。
まずカレントディレクトリでファイルを検索し、見つからない場合はスクリプトが
配置されているディレクトリで検索します。
:param infile: str, optional: 検索対象のファイル名。デフォルトはNone。
:returns: str or None: 見つかったファイルの絶対パス、またはNone。
"""
script_path = os.path.abspath(sys.argv[0])
script_dir = os.path.dirname(script_path)
script_name = os.path.splitext(os.path.basename(script_path))[0]
default_ini = f"{script_name}.ini"
search_target = infile if infile else default_ini
candidate1 = os.path.join(os.getcwd(), search_target)
if os.path.isfile(candidate1):
return candidate1
candidate2 = os.path.join(script_dir, search_target)
if os.path.isfile(candidate2):
return candidate2
return None
[ドキュメント]
def read_ini(inifile=None):
"""
INIファイルの内容を読み込み、辞書として返す。
`search_file` を使用してINIファイルのパスを特定し、ファイルを読み込みます。
`#` または `;` で始まる行はコメントとしてスキップされます。
`key = value` 形式の行を解析し、トリプルクォート (`\"\"\"` または `\'\'\'`) を
使用した複数行の値をサポートします。また、読み込み後には `$VAR` 形式の
変数を展開します。
:param inifile: str, optional: 読み込むINIファイルのパス。デフォルトはNone。
:returns: dict: INIファイルから読み込んだ設定を格納する辞書。
:raises FileNotFoundError: 指定されたINIファイルが見つからない場合。
"""
path = search_file(inifile)
if path is None:
raise FileNotFoundError(f"INIファイルが見つかりませんでした: {inifile}")
result = {}
variables = {}
current_key = None
multiline_val = []
multiline_delim = None
with open(path, 'r', encoding='utf-8') as f:
for line in f:
line = line.rstrip()
if not line or line.startswith('#') or line.startswith(';'):
continue
if multiline_delim:
if line.strip() == multiline_delim:
val = '\n'.join(multiline_val)
result[current_key] = val
variables[current_key] = val
current_key = None
multiline_val = []
multiline_delim = None
else:
multiline_val.append(line)
continue
if '=' in line:
key, val = map(str.strip, line.split('=', 1))
if (val == '\"\"\"' or val == "\'\'\'" or
(val.startswith('\"\"\"') and not val.endswith('\"\"\"')) or
(val.startswith("\'\'\'") and not val.endswith("\'\'\'"))):
multiline_delim = val[:3]
content = val[3:]
multiline_val = [content] if content else []
current_key = key
continue
if (val.startswith('\"\"\"') and val.endswith('\"\"\"')) or \
(val.startswith("\'\'\'") and val.endswith("\'\'\'")):
val = val[3:-3]
result[key] = val
variables[key] = val
for key, val in result.items():
def expand_var(match):
return variables.get(match.group(1), match.group(0))
result[key] = re.sub(r"\$(\w+)\b", expand_var, val)
return result
# =========================================================
# メイン処理
# =========================================================
[ドキュメント]
def initialize():
"""
コマンドライン引数パーサーを初期化し、設定する。
`argparse.ArgumentParser` を初期化し、ツールの説明を設定します。
必要な引数（`pattern`, `output`）とオプション引数（`--inifile`, `--api`,
`--update`, `--overwrite`, `--pause`）を追加します。
`default_ini_name` は実行スクリプト名に基づいて生成されます。
:returns: argparse.ArgumentParser: 設定済みのパーサーオブジェクト。
"""
default_ini_name = os.path.splitext(os.path.basename(sys.argv[0]))[0] + ".ini"
parser = argparse.ArgumentParser(description="PythonコードにSphinx形式のDocstringを自動追加するツール")
# explain_program5.py に合わせた引数構成
parser.add_argument("pattern", help="対象ファイルのワイルドカード（例: '*.py'）")
parser.add_argument("output", nargs="?", help="出力名（単一ファイル時のみ）")
parser.add_argument("--inifile", default=default_ini_name, help="プロンプト設定ファイルパス")
parser.add_argument("--api", choices=["openai", "openai5", "google", "gemini"], default='google')
parser.add_argument("-u", "--update", type=int, default=0, help="1の場合、ソースが新しい場合のみ更新")
parser.add_argument("-w", "--overwrite", type=int, default=0, help="1の場合、既存ファイルを上書き")
parser.add_argument("-p", "--pause", type=int, default=1)
return parser
[ドキュメント]
def read_args(parser):
"""
ArgumentParserから引数を解析し、その値を表示する。
`parser.parse_args()` を呼び出して引数を解析し、
解析された引数の値を標準出力に表示します。
:param parser: argparse.ArgumentParser: 初期化済みのArgumentParserオブジェクト。
:returns: argparse.Namespace: 解析された引数を格納するオブジェクト。
"""
args = parser.parse_args()
print("Args:")
print(f" {args.pattern=}")
print(f" {args.output=}")
print(f" {args.inifile=}")
print(f" {args.api=}")
print(f" {args.update=}")
print(f" {args.overwrite=}")
print(f" {args.pause=}")
return args
[ドキュメント]
def main():
"""
スクリプトのメイン処理を実行する。
`initialize` と `read_args` を呼び出してコマンドライン引数を処理します。
AI設定ファイル (`ai.env`) を読み込み、指定されたINIファイルからプロンプト設定を取得します。
ワイルドカードパターンに合致するファイル群を処理し、各ファイルについて
更新・上書きルールに従って処理をスキップするか判断します。
ファイル内容を読み込み、プロンプトを構築してAIモデルを呼び出します。
AIからの応答からDocstringを抽出し、指定された出力ファイルに書き込みます。
エラー発生時はメッセージを表示し、スタックトレースを出力します。
API呼び出しの負荷軽減のため、ファイル処理ごとに1秒待機し、
`pause` 引数が設定されている場合、終了前にユーザー入力を待ちます。
:returns: None
:raises SystemExit: INIファイルの読み込み失敗時や、マッチするファイルがない場合に発生。
"""
print()
print(f"=== {sys.argv[0]} ===")
parser = initialize()
args = read_args(parser)
read_ai_config("ai.env")
try:
ini_data = read_ini(args.inifile)
print(f"Loaded INI: {search_file(args.inifile)}")
except Exception as e:
print(f"Error loading INI: {e}")
sys.exit(1)
# ワイルドカード展開
files = glob.glob(args.pattern)
if not files:
print(f"No files matched pattern: {args.pattern}")
sys.exit(1)
# 出力ファイル名のリスト作成
if args.output and len(files) == 1:
outputs = [args.output]
else:
outputs = [os.path.splitext(f)[0] + "_docstring.py" for f in files]
for inp, out in zip(files, outputs):
# 更新・上書きチェックロジック
if os.path.exists(out) and not args.overwrite:
# updateモードかつ、出力ファイルの方が新しい場合はスキップ
if not args.update or os.path.getmtime(out) >= os.path.getmtime(inp):
print(f"Skip: {out}")
continue
print(f"Processing: {inp} -> {out}")
try:
code = Path(inp).read_text(encoding="utf-8")
except:
try:
code = Path(inp).read_text(encoding="shift-jis")
except Exception as e:
print(f"Error reading {inp}: {e}")
continue
base_name = os.path.splitext(os.path.basename(inp))[0]
template = ini_data.get("PROMPT_MAIN", "")
role = ini_data.get("SYSTEM_ROLE", "Assistant")
prompt = template.replace("{{script_name}}", inp)\
.replace("{{code}}", code)\
.replace("add_docstring", base_name)
try:
if args.api == "openai5":
res = query_openai5(prompt, os.getenv("openai_model5"), instructions=role)
doc = extract_openai5_text(res)
elif args.api == "openai":
res = query_openai4(prompt, os.getenv("openai_model"), role=role)
doc = res.choices[0].message.content
else:
res = query_google(prompt, os.getenv("gemini_model"), role=role)
doc = res.text if res else None
if doc:
# Markdownブロックの除去
doc = doc.strip()
if doc.startswith("```"):
doc = re.sub(r'^```[a-zA-Z]*\n', '', doc)
doc = re.sub(r'\n```$', '', doc)
Path(out).write_text(doc, encoding="utf-8")
print(f"Done: {out}")
else:
print(f"Error: No response from AI for {inp}")
except Exception:
print(f"\n!!! Error during AI processing for {inp} !!!")
traceback.print_exc()
time.sleep(1) # API連続呼び出しの負荷軽減
if args.pause:
input("\nPress ENTER to terminate>>\n")
if __name__ == "__main__":
main()