add_docstring.py の技術ドキュメント

プログラムの動作

add_docstring.py は、指定されたPythonソースコードファイルに対して、AI(OpenAIまたはGoogle Gemini)を利用してSphinx形式のDocstringを自動的に追加するツールです。主に以下の機能を提供します。

  • Docstringの自動生成: 入力されたPythonコードの内容をAIに渡し、適切なDocstring(関数、クラス、モジュールレベル)を生成させます。生成されるDocstringはSphinxの形式に準拠します。

  • 柔軟なファイル指定: ワイルドカードパターン (*.py) を用いて複数のファイルを一度に処理できます。

  • INIファイルによる設定: プロンプトテンプレート、AIモデル、システムロールなどの設定をINIファイルから読み込むことで、AIの動作を柔軟にカスタマイズできます。INIファイルは複数行の記述や変数展開に対応しています。

  • 更新・上書き制御: 既存の出力ファイルが存在する場合、上書きオプション (--overwrite) や、入力ファイルより出力ファイルが新しい場合に処理をスキップする更新オプション (--update) を利用して、不必要なAI呼び出しやファイル変更を防ぎます。

  • AIサービス連携: tkai_lib を通じて、OpenAIのGPTシリーズ(GPT-4、GPT-3.5など)およびGoogle Gemini(旧Bard)APIと連携します。

このプログラムは、Pythonプロジェクトにおけるドキュメンテーション作業の効率化、特にDocstringの記述にかかる時間と労力を削減することを目的としています。

原理

add_docstring.py は、主に以下のアルゴリズムと仕組みに基づいて動作します。

  1. 初期化と引数解析: argparse モジュールを使用してコマンドライン引数を解析します。入力ファイルのワイルドカードパターン、出力ファイル名、INIファイルパス、使用するAIサービス、更新・上書きモードなどが設定されます。

  2. INIファイルの読み込み:

    • search_file 関数は、指定されたINIファイル(デフォルトはスクリプト名.ini)をカレントディレクトリまたはスクリプト実行ディレクトリから検索します。

    • read_ini 関数は、INIファイルを解析し、設定値を辞書として読み込みます。この関数は以下の特徴を持ちます。

      • コメント行 (# または ; で始まる行) を無視します。

      • key = value 形式で値を読み込みます。

      • 値が 3重引用符 で始まる場合、複数行の値をサポートし、対応する 3重引用符 が現れるまでを行を結合して1つの値とします。

      • $VAR 形式の変数参照をサポートし、INIファイル内で定義された他の変数に展開します。

  3. AI連携設定の読み込み: tkai_lib.read_ai_config("ai.env") を呼び出し、ai.env ファイルからAIサービスへの接続に必要な環境変数(APIキー、モデル名など)をロードします。

  4. 対象ファイルの特定と出力ファイル名の決定:

    • glob.glob を使用して、コマンドライン引数で指定されたワイルドカードパターンに一致するPythonソースファイル群を特定します。

    • 出力ファイル名は、単一ファイルかつ --output 引数が指定されていればその値を使用します。それ以外の場合は、元のファイル名に _docstring.py というサフィックスを追加した名前とします。

  5. ファイル処理ループ: 検出された各入力ファイルに対して以下の処理を繰り返します。

    • 更新・上書きチェック:

      • 出力ファイルが既に存在し、--overwrite オプションが指定されていない場合、処理をスキップする可能性があります。

      • --update オプションが有効な場合、出力ファイルの最終更新時刻が入力ファイルの最終更新時刻よりも新しい、または同じであれば処理をスキップします。これにより、AIへの不要なリクエストを避けます。

      • 最終更新時刻の比較には os.path.getmtime() を使用します。

    • ソースコードの読み込み: 入力Pythonファイルを読み込みます。エンコーディングはUTF-8を優先し、失敗した場合はShift-JISを試行します。

    • プロンプトの生成:

      • INIファイルから読み込んだ PROMPT_MAIN テンプレートに、入力ファイル名 ({{script_name}})、ソースコード内容 ({{code}})、ベース名 ({{base_name}}) を埋め込み、AIへのリクエストプロンプトを構築します。

      • INIファイルで定義された SYSTEM_ROLE をAIのシステム指示として使用します。

    • AIの呼び出し:

      • --api 引数に応じて、tkai_libquery_openai4, query_openai5, query_google のいずれかの関数を呼び出し、AIにプロンプトを送信します。

      • AIからの応答を待ち受けます。

    • AI応答の処理と保存:

      • AIから返されたテキストから、もし存在すればMarkdownのコードブロックを示す 3重引用符 と言語指定(例: ```python`)の開始/終了記号を除去します。

      • 整形されたDocstring付きコードを、決定された出力ファイル名でUTF-8エンコーディングで保存します。

  6. レートリミット対策: 各AI呼び出しの間に time.sleep(1) を挿入し、APIへの連続リクエストによるレートリミット超過を防ぎます。

必要な非標準ライブラリとインストール方法

add_docstring.py は、AI連携のために tkai_lib というカスタムライブラリと、それが依存する公式のAIクライアントライブラリを必要とします。

  1. tkai_lib: このライブラリは、配布されるプログラムの一部として提供されるか、ユーザーが別途準備する必要があります。add_docstring.py と同じディレクトリ、またはPythonのパスが通った場所に tkai_lib.py ファイルを配置してください。

  2. openai: OpenAI APIを使用する場合に必要です。

    pip install openai
    
  3. google-generativeai: Google Gemini APIを使用する場合に必要です。

    pip install google-generativeai
    

これらのライブラリがインストールされていない、または tkai_lib.py が見つからない場合、プログラムはエラーメッセージを出力して終了します。

必要な入力ファイル

add_docstring.py は、以下のファイルを入力として利用します。

  1. Pythonソースコードファイル: Docstringを追加したいPythonの .py ファイル。ワイルドカードパターン (*.py) を使用して複数指定できます。

  2. INI設定ファイル: AI連携のためのプロンプトや設定を記述するファイルです。

    • デフォルト名: add_docstring.ini (スクリプト名に基づきます)。

    • 指定方法: --inifile <ファイルパス> オプションで明示的に指定できます。

    • 検索順序: まずカレントディレクトリ、次に add_docstring.py が存在するディレクトリを検索します。

    • 主要な設定項目:

      • PROMPT_MAIN: AIに与えるプロンプトのテンプレート。以下のプレースホルダーが利用可能です。

        • {{script_name}}: 処理中のスクリプトのファイル名。

        • {{code}}: 処理中のスクリプトの全コード内容。

        • {{base_name}}: 処理中のスクリプトの拡張子を除いたベース名。

      • SYSTEM_ROLE: AIに対するシステムとしての役割指示。

    INIファイルの例 (add_docstring.ini):

    # AI連携プロンプト設定
    SYSTEM_ROLE = You are a helpful AI assistant that specializes in generating Sphinx-style docstrings for Python code.
    PROMPT_MAIN = '''
    Please analyze the following Python script and add Sphinx-style docstrings to all functions, classes, and the module itself.
    Ensure the docstrings are comprehensive, explaining parameters, return values, and exceptions where applicable.
    Provide the complete Python code with the added docstrings. Do not include any extra text or explanations.
    
    Script name: {{script_name}}
    
    Python Code:
    ```python
    {{code}}
    

    '''

    上記例のように、`PROMPT_MAIN` のような長い値は `3重引用符` で囲むことで複数行にわたって記述できます。
    
    
  3. 環境変数設定ファイル (ai.env): AIサービスのAPIキーや使用するモデル名を環境変数として設定するファイルです。tkai_lib が読み込みます。ファイル名は固定で ai.env です。

    • OpenAIの場合:

      openai_api_key=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
      openai_model=gpt-4
      openai_model5=gpt-3.5-turbo
      
    • Google Geminiの場合:

      google_api_key=AIzaSyxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
      gemini_model=gemini-pro
      

    これらのキーは利用するAIサービスによって異なります。

生成される出力ファイル

プログラムの実行により、Docstringが追加された新しいPythonソースコードファイルが生成されます。

  • ファイル名:

    • ワイルドカードで複数のファイルを指定した場合、または単一ファイルでも --output オプションが指定されていない場合: 元のファイル名の末尾に _docstring.py が追加されます。 例: my_script.pymy_script_docstring.py

    • 単一ファイルを指定し、かつ --output <出力ファイル名> オプションが指定されている場合: 指定された <出力ファイル名> で保存されます。

  • 内容: 入力されたPythonコードに、AIが生成したSphinx形式のDocstringが追加された完全なPythonコード。

コマンドラインでの使用例 (Usage)

基本的なコマンドラインの書式は以下の通りです。

python add_docstring.py <pattern> [output] [--inifile INIFILE] [--api {openai,openai5,google,gemini}] [-u UPDATE] [-w OVERWRITE] [-p PAUSE]
  • <pattern>: 対象とするPythonファイルのワイルドカードパターン(例: '*.py', 'my_module.py')。必須。

  • [output]: 出力ファイル名。単一の入力ファイルのみを指定した場合に有効です。省略した場合、入力ファイル名に _docstring.py が付加されます。

  • --inifile INIFILE: プロンプト設定が記述されたINIファイルのパス。デフォルトは add_docstring.ini

  • --api {openai,openai5,google,gemini}: 使用するAIサービスを指定します。

    • openai: OpenAIのGPT-4モデル(openai_model 環境変数で指定)。

    • openai5: OpenAIのGPT-3.5モデル(openai_model5 環境変数で指定)。

    • google: Google Geminiモデル(gemini_model 環境変数で指定)。

    • gemini: google と同じ。 デフォルトは google

  • -u UPDATE, --update UPDATE: 1 を指定すると、出力ファイルが入力ファイルよりも新しい場合に処理をスキップします。デフォルトは 0 (無効)。

  • -w OVERWRITE, --overwrite OVERWRITE: 1 を指定すると、既存の出力ファイルを無条件に上書きします。デフォルトは 0 (無効)。

  • -p PAUSE, --pause PAUSE: 1 を指定すると、すべての処理後に続行するためのEnterキー入力を求めます。デフォルトは 1 (有効)。

コマンドラインでの具体的な使用例

以下のファイルが準備されていると仮定します。

  1. sample_code.py (入力Pythonファイル)

    def calculate_area(width, height):
        return width * height
    
    class Rectangle:
        def __init__(self, width, height):
            self.width = width
            self.height = height
    
        def get_area(self):
            return self.width * self.height
    
  2. config.ini (INI設定ファイル)

    SYSTEM_ROLE = You are a helpful AI assistant that specializes in generating Sphinx-style docstrings for Python code.
    PROMPT_MAIN = '''
    Please analyze the following Python script and add comprehensive Sphinx-style docstrings to all functions, classes, and the module itself.
    Ensure the docstrings explain parameters, return values, and any exceptions where applicable.
    Provide the complete Python code with the added docstrings. Do not include any extra text or explanations outside the code block.
    
    Python Code:
    ```python
    {{code}}
    

    '''

    
    
  3. ai.env (環境変数設定ファイル)

    google_api_key=AIzaSy_YOUR_GOOGLE_API_KEY_HERE
    gemini_model=gemini-pro
    

    AIzaSy_YOUR_GOOGLE_API_KEY_HERE の部分は、実際に取得したGoogle Gemini APIキーに置き換えてください。

実行コマンド例

sample_code.py にDocstringを追加し、sample_code_docstring.py として保存する場合。

python add_docstring.py sample_code.py --inifile config.ini --api google

実行結果の説明

上記のコマンドを実行すると、以下のようなメッセージがコンソールに出力され、sample_code_docstring.py という新しいファイルが生成されます。

=== add_docstring.py ===
Args:
  args.pattern='sample_code.py'
  args.output=None
  args.inifile='config.ini'
  args.api='google'
  args.update=0
  args.overwrite=0
  args.pause=1
Loaded INI: /path/to/your/config.ini
Processing: sample_code.py -> sample_code_docstring.py
Done: sample_code_docstring.py

Press ENTER to terminate>>

生成された sample_code_docstring.py の内容は、sample_code.py の各関数とクラスにSphinx形式のDocstringがAIによって追加されたものになります(Docstringの内容はAIの生成結果により変動します)。

例: 生成される sample_code_docstring.py の内容(AIの出力例)

"""
This module provides utilities for geometric calculations, specifically for rectangles.
"""

def calculate_area(width, height):
    """Calculate the area of a rectangle.

    :param width: The width of the rectangle.
    :type width: float or int
    :param height: The height of the rectangle.
    :type height: float or int
    :return: The area of the rectangle.
    :rtype: float or int
    """
    return width * height

class Rectangle:
    """Represents a rectangle with given width and height.

    :param width: The width of the rectangle.
    :type width: float or int
    :param height: The height of the rectangle.
    :type height: float or int
    """
    def __init__(self, width, height):
        self.width = width
        self.height = height

    def get_area(self):
        """Get the area of the rectangle.

        :return: The area of the rectangle.
        :rtype: float or int
        """
        return self.width * self.height