Python + astモジュールを使ってソースコードを解析し、メソッドブロックや関数ブロックの定義行と最終行を取得する

Python

Pythonソースコードを解析して、メソッドブロックや関数ブロックの定義行と最終行を取得することがありました。

Pythonでは標準モジュールのastを使ってソースコードを解析できるため、試した時のメモです。
32.2. ast — 抽象構文木 — Python 3.6.1 ドキュメント

なお、以下のページが大変参考になりました。ありがとうございました。
Python: ast (Abstract Syntax Tree: 抽象構文木) モジュールについて - CUBE SUGAR CONTAINER

　

環境

　

方法

上記の参考ブログ同様、astモジュールと再帰を組み合わせて解析します。

  • メソッドや関数の定義は、ast.FunctionDefクラス
  • メソッド名や関数名は、ast.FunctionDef.nameで取得
  • 行数はast.FunctionDef.linenoなどのlineno属性で取得
    • クラスによっては、lineno属性がないことに注意

　
collect_method_last_line_no.py

import ast

class Collection:
    def __init__(self, name='', def_line_no=0, last_line_no=0):
        self.name = name
        self.def_line_no = def_line_no
        self.last_line_no = last_line_no


class MethodLastLineNoCollector:
    def __init__(self):
        self.result = {}
        self.searched_line_no = 0

    def run(self, node):
        if isinstance(node, ast.FunctionDef):
            self.result[node.lineno] = Collection(node.name, def_line_no=node.lineno)

        for child in ast.iter_child_nodes(node):
            self.run(child)
        
        if hasattr(node, 'lineno'):
            # 探索した最終行を取得
            # 再帰で探すので、node.linenoは 1 > 2 > 3 > 2 > 1 となる
            if node.lineno > self.searched_line_no:
                self.searched_line_no = node.lineno

            # 再帰で探した時の帰りに、最終行を設定する
            else:
                if self.result.get(node.lineno):
                    self.result[node.lineno].last_line_no = self.searched_line_no


if __name__ == '__main__':
    FILENAME = 'target.py'
    with open(FILENAME, 'r') as f:
        source = f.read()

    tree = ast.parse(source, FILENAME)

    collector = MethodLastLineNoCollector()
    collector.run(tree)

    for v in collector.result.values():
        print(f'{v.name} -> def:{v.def_line_no}, last:{v.last_line_no}')

　
動作確認をします。例えば、

target.py

def foo_function():
    pass

    def innner_foo_function():
        pass


class Bar:
    def bar_method(self):
        pass

        def bar_inner_method(self):
            pass

    def bar_other_method(self):
        pass

    class InnerBar:
        def innter_bar_method(self):
            pass


if __name__ == "__main__":
    pass

というソースコードがあったとします。

　
collect_method_last_line_no.pyを実行します。

$ python collect_method_last_line_no.py 
foo_function -> def:1, last:5
innner_foo_function -> def:4, last:5
bar_method -> def:9, last:13
bar_inner_method -> def:12, last:13
bar_other_method -> def:15, last:16
innter_bar_method -> def:19, last:20

動作しているようです。

　

ソースコード

GitHubに上げました。collect_method_last_line_noディレクトリの中が今回のものです。
thinkAmi-sandbox/python_ast-sample