「Python でデコレートされているメソッドを取得する方法」をやってみた

元ネタ: Python でデコレートされているメソッドを取得する方法

元ブログでは inspect モジュールで調べていますが、こういうのは ast モジュールを使うと便利ですね。

find_decorator.py

# -*- coding: utf-8 -*-
import ast

source = """
class Sample(object):

    @foo_decorator2
    @foo_decorator1
    def foo(self):
        pass

    @bar_decorator
    def bar(self):
        pass

    def baz(self):
        pass

@hoge_decorator2
@hoge_decorator1
def hoge(x):
    pass

def fuga(x):
    pass
"""

class FindDecorator(ast.NodeVisitor):
    def visit_FunctionDef(self, node):
        decorators = node.decorator_list
        if decorators:
            print('decorated: {}'.format(node.name))
            for name in decorators:
                print('decorator: {}'.format(name.id))
            print('-' * 32)
        return node

FindDecorator().visit(ast.parse(source))

こんな感じ。

$ python find_decorator.py 
decorated: foo
decorator: foo_decorator2
decorator: foo_decorator1
--------------------------------
decorated: bar
decorator: bar_decorator
--------------------------------
decorated: hoge
decorator: hoge_decorator2
decorator: hoge_decorator1
--------------------------------

追記:
inspect.getsource を使って、モジュールのソースコードを動的に取得することもできました。

3.4

# -*- coding: utf-8 -*-
import ast
import inspect
import sys

def foo_decorator2(func): pass
def foo_decorator1(func): pass
def bar_decorator(func): pass
def hoge_decorator2(func): pass
def hoge_decorator1(func): pass

class Sample(object):

    @foo_decorator2
    @foo_decorator1
    def foo(self):
        pass

    @bar_decorator
    def bar(self):
        pass

    def baz(self):
        pass

@hoge_decorator2
@hoge_decorator1
def hoge(x):
    pass

def fuga(x):
    pass

class FindDecorator(ast.NodeVisitor):
    def visit_FunctionDef(self, node):
        decorators = node.decorator_list
        if decorators:
            print('decorated: {}'.format(node.name))
            for name in decorators:
                print('decorator: {}'.format(name.id))
            print('-' * 32)
        return node

module = sys.modules[__name__]
source = inspect.getsource(module)
FindDecorator().visit(ast.parse(source))

inspect.getsource の引数に渡したオブジェクトのソースコードを取得できる。ここではモジュールそのものを渡しているけど、クラスや関数なども渡すことができる。