Diary?

(18:46)

「デコレートのタイミングで、デコレート対象がトップレベルの関数であるかクラスのメソッドであるかを判別する方法はあるだろうか」という問題が出たんで「確か Python では関数とメソッドは types で識別できたよなー」と思ったんで、実際にやってみた（以下、 bpython の画面のコピペ）。

>>> class FOo(object):
...     def foo(self):
...         pass
>>> def bar():
...     pass
>>> import types
>>> isinstance(Foo().foo, types.MethodType)
True
>>> isinstance(bar, types.FunctionType)
True

これをデコレータでやるとこんな感じか。

>>> def dec(f):
...     def _(*args, **kwds):
...         if isinstance(f, types.MethodType):
...             print 'method'
...         elif isinstance(f, types.FunctionType):
...             print 'function'
...         return f(*args, **kwds)
...     return _
...     
... 
>>> class Hoge(object):
...     @dec
...     def hoge(self):
...         print 'hoge'
... 
>>> @dec
... def hage():
...     print 'hage'
... 
>>> Hoge().hoge()
function
hoge

「ええ、何で!?」と一瞬思うけど、よくよく考えれば dec に関数が渡された時点では、それはまだ特定のクラスのインスタンスに属しているわけではないのだ。その辺は dis モジュールで逆アセンブルしつつ、オフィシャルのバイトコード命令を見てればなんとなく掴めるはず。

じゃあどうすんのって話だけど、実は俺は昔に冒頭の課題のようなことをやろうとした事があって、そんときにやったやり方は次のようなものだった。と思う。

1. inspect モジュールでスタックフレームを取り出す
2. デコレータを適用してる場所を調べる。デコレータ名を LOAD_NAME しているところがそれ
3. STORE_NAME でデコレータの適用されたオブジェクトの名前がわかるので、それを取得する
4. 一連のスタックフレームを実行してる関数名（これはクラス名の事もある）を調べる
5. その関数名を STORE_NAME してる部分をしらべ、その直前で BUILD_CLASS していればクラスのメソッドへのデコレータ

実にバカバカしいというか無理やりなやり方だが、一応できなくもない。真面目にやるならもっといろんなパターンを試す必要があるが（例えば STORE_NAME の他にも STORE 系命令あるし）、まあ大雑把にやるとこうだろう。

import dis
import re
import sys
from StringIO import StringIO
class DecoratorInspector(object):
    def __init__(self, name, stackframelist):
        orig = sys.stdout
        self.name = name
        self.codes = {}
        try:
            for s in stackframelist:
                sys.stdout = StringIO()
                dis.dis(s[0].f_code)
                sys.stdout.seek(0)
                self.codes[s[3]] = sys.stdout.readlines()
        finally:
            sys.stdout = orig
    
    def list_decorating_point(self):
        for defname, lines in self.codes.iteritems():
            for i, line in enumerate(lines):
                if '(%s)' % self.name in line and 'LOAD_NAME':
                    yield defname, i
    
    def list_decorated_function(self, defname, index):
        for line in self.codes[defname][index:]:
            if 'STORE_NAME' in line:
                fname = line.split('(')[1].split(')')[0]
                yield defname, fname
    
    def isDefinedInClass(self, defname):
        for lines in self.codes.itervalues():
            for i, line in enumerate(lines):
                if '(%s)' % defname in  line and ' STORE_NAME ' in line:
                    return 'BUILD_CLASS' in lines[i - 1]
    
    def isMethod(self, func):
        defmap = {}
        for n, i in self.list_decorating_point():
            for d, f in self.list_decorated_function(n, i):
                defmap[f] = d
        n = func.__name__
        return self.isDefinedInClass(defmap[n])

使い方はこう。

import inspect
def dec(f):
    di = DecoratorInspector('dec', inspect.stack())
    if di.isMethod(f):
        print '%s is method' % (f.__name__)
    else:
        print '%s is function' % (f.__name__)
    return f
    
class Foo(object):
    @dec
    def foo(self):
        print 'foo!'
    
    def bar(self):
        print 'bar!'
    
@dec
def buz():
    print 'buz!'

このコードを実行すると、次のような出力が得られる。

foo is method
buz is function

かなり穴のありそうな気がするが、俺の知ってる限りじゃこういうやり方しか思いつかなかった。誰かもっとマシなやり方しらないかなあ。