Trieではてなキーワード自動リンク

pytstを使ってはてなキーワード自動リンク。全部 Unicode でやりたかったけど、pytst が unicode オブジェクトを扱えないらしくとりあえず、内部は全部 UTF-8 の str で。
例のdartsを使ったやつと違って一応大文字小文字を区別しないマッチングをするようにしてある。
ちなみにこの間抽出したのを使って 400KBytes くらいのファイルに試したら Perl ではてなキーワード自動リンクAPIの正規表現版をそのまま使ったやつの10倍くらい早かった。

import urllib
import tst

class KeywordAutoLinkCallback:
    def __init__(self, original_text):
        self.buffer = ""
        self.original_text = original_text
        self.offset = 0

    def match(self, key, length, obj):
        text = self.original_text[self.offset:self.offset + len(key)]
        self.offset += len(key)
        if length > 0:
            self.buffer += self.make_keyword_link(text)
        else:
            self.buffer += text

    def result(self):
        return self.buffer

    def make_keyword_link(self, keyword):
        return '<a href="%s">%s</a>' % (self.make_keyword_url(keyword), keyword)

    def make_keyword_url(self, keyword):
        return 'http://d.hatena.ne.jp/keyword/' + urllib.quote(keyword)

if __name__ == '__main__':
    import sys, locale, fileinput

    if len(sys.argv) < 2:
        print >>sys.stderr, "usage: %s keywordlist.extracted [doc1] [doc2] ... " % sys.argv[0]
        sys.exit(1)

    internal_encoding = 'utf-8'
    io_encoding = locale.getpreferredencoding()
    keyword_filename = sys.argv[1]
    doc_filenames = sys.argv[2:]

    trie = tst.TST()
    for line in open(keyword_filename, 'r'):
        keyword = line.rstrip()
        trie[keyword] = 1

    def conv(line):
        return line.decode(io_encoding).rstrip().encode(internal_encoding)

    for line in (conv(line) for line in fileinput.input(doc_filenames)):
        c = KeywordAutoLinkCallback(line)
        result = trie.scan(line.lower(), tst.CallableAction(c.match, c.result))
        print result.decode(internal_encoding).encode(io_encoding)