Pythonでデータ定義行のファイル名と行番号を手に入れる

ふと思いついたので書いておく。

Pythonでデータ定義行のファイル名と行番号を手に入れる - 清水川Web では、Pythonでデータの定義位置を記録する方法として、データ生成用の関数を作ってその中でデータの定義位置を記録する方式が提案されている。

このようにデータの定義位置を参照したい、というケースはたまにあるが、上記の方式はあまり汎用的ではない。登録できるデータとして辞書が必要だし、エラー発生時のメッセージを自分で生成するのは面倒だ。いちいち関数呼び出し形式で書かなければならないのもあまり好きではない。

で、ちょっと考えてみたのがこれだ。

import sys

class _DataLoggerItem:
    def __init__(self, logger):
        self._logger = logger
        self._iter = iter(self._logger._data)
        
    def next(self):
        ret = self._iter.next()
        self._logger._curvalue = ret
        return ret

class DataLogger(object):
    _obj_none = object()
    
    def __init__(self):
        self._data = []
        self._loc = {}
        self._curvalue = self._obj_none
        
    def __lshift__(self, other):
        self._data.append(other)

        frame = sys._getframe().f_back
        self._loc[id(other)] = frame.f_code.co_filename, frame.f_lineno
    
        return self

    def __iter__(self):
        return _DataLoggerItem(self)
    
    def __enter__(self):
        return self
    
    def __exit__(self, type, value, traceback):
        if type and self._curvalue is not self._obj_none:
            print "Error at %s:%d" % (self._loc[id(self._curvalue)])

などとしておいて、テストコードなどでは

data = DataLogger()
(data 
<< 1  
<< 2  
<< 3  
<< 4 )

with data:
    for d in data:
        if d == 3:
            raise RuntimeError

などと書く。<< 演算子を悪用して関数呼び出しを回避し、context managerでエラーメッセージを表示するアイデアだ。これならもうちょっと使いやすいのではないだろうか。