このエントリでは、
Pythonを使って、Pythonの関数のソースコードを抜き出す方法を示します。

日本語で説明しても、
何をしようとしているのか分かりにくいと思うので、
まず、コードと実行結果を示します。

コード: get_func_code.py

import ast
from typing import List

def get_func_code(f) -> str:
    """
    指定した関数のソースコードを返却する
    """
    def walk(nodes, lineno, hit):
        endlineno = -1
        for i, n in enumerate(nodes):
            if hit:
                endlineno = n.lineno
            elif isinstance(n, ast.FunctionDef) and n.lineno == lineno:
                hit = True
            else:
                endlineno, hit = walk(ast.iter_child_nodes(n), lineno, hit)
            if endlineno != -1:
                break
        return endlineno, hit
    filename = f.__code__.co_filename
    lineno = f.__code__.co_firstlineno
    with open(filename) as f:
        source = f.read()
    tree = ast.parse(source, filename)
    endlineno, hit = walk(tree.body, lineno, False)
    if endlineno == -1:
        return '\n'.join(source.split('\n')[lineno - 1:])
    return '\n'.join(source.split('\n')[lineno - 1:endlineno - 1])

if __name__ == '__main__':
    code = get_func_code(ast.parse)
    print(code)

実行結果:

% python3 get_func_code.py
def parse(source, filename='<unknown>', mode='exec', *,
          type_comments=False, feature_version=None):
    """
    Parse the source into an AST node.
    Equivalent to compile(source, filename, mode, PyCF_ONLY_AST).
    Pass type_comments=True to get back type comments where the syntax allows.
    """
    flags = PyCF_ONLY_AST
    if type_comments:
        flags |= PyCF_TYPE_COMMENTS
    if isinstance(feature_version, tuple):
        major, minor = feature_version  # Should be a 2-tuple.
        assert major == 3
        feature_version = minor
    elif feature_version is None:
        feature_version = -1
    # Else it should be an int giving the minor version for 3.x.
    return compile(source, filename, mode, flags,
                   _feature_version=feature_version)

上記のコードを実行すると、
ast.parse関数のソースコードが表示されます。

このようにして、Pythonを使って、
Pythonの関数のソースコードを抜き出すことが出来ました。

get_func_code関数で、関数を受け取りその実装を返却します。
ざっくりとした処理の流れは、次のようになっています。

  • 関数のオブジェクトから、実装のファイル名・関数の開始位置をとる
  • 実装のファイルを開き、構文木を使って、関数の終了位置(次の定義の開始位置)を調べる
  • 実装のファイルの、開始位置〜終了位置の文字列を返却

ちなみに、どうして「関数の実装を抜き出したい」と思ったのかと言うと、
関数の実装をチェックする単体テストを書きたかったからです。

利用するFramework等によっては、
決まったメソッドに定型的なコードを記載する必要があったりします。
ですが、この手の機械的なコードはミスしがちなので、
単体テストでチェック出来ないかと思い、この方法を考えました。

Frameworkに合わせたLinterとかを作るのが正攻法かも知れないですが、
単体テストで対処するのもお手軽で悪くないかなと思っています。