difflib の使い方 〜hg-diff-highlight を例にして〜

hg-diff-highlight という Mercurial 拡張を作った話を書いたのですが、
このエントリではその実装について、特に内部で使っている difflib の使い方についてご紹介します。

hg との連携部分

hg-diff-highlight では mercurial の diff 計算部分には手を入れていません。
mercurial では mercurial.patch パッケージに diff 関連のロジックがいろいろ入っているようなのですが、
今回はハイライトだけに注力するため特に手を入れていません。

代わりに手を入れたのが diff 出力部分です。
mercurial の diff はすべて mercurial.ui パッケージの ui クラスを介して出力されています。
color 拡張を使うと、ui クラスの代わりに hgext.color パッケージの colorui クラスが利用されます。
hg-diff-highlight ではこの ui インターフェースをさらに拡張し、diff 出力の文字単位でのハイライトを行うようにしています。

diff 計算部分では ui#write() メソッドを使って diff を書きだすのですが、その際にカラーリング用のラベル情報を付記しています。
たとえば、hello world する python コードを書き換えると次のように ui#write() メソッドが呼び出されます。

ui.write('def say():', label='')
ui.write('\n', label='')
ui.write('-    print "hello WORLD"', label='diff.deleted')
ui.write('\n', label='')
ui.write('+    print "hello world"', label='diff.inserted')
ui.write('\n', label='')

color 拡張はこのラベルを見てテキストのカラーリングをしているというわけですね。
Mercurial ではテキストを出力する部分は基本的に ui モジュールを介しているようなので、
ui モジュールを拡張すると入出力周りをハックしやすくなります。

ちなみに、この ui#write() ですが、基本的には 1行につき 2回(テキスト部分と改行コード) ずつ呼び出されるのですが、
差分がある行に trailing whitespace (行末のスペース)がある場合は、その部分だけ別ラベルで呼び出されます。

ui.write('def say():', label='')
ui.write('\n', label='')
ui.write('-    print "hello WORLD"', label='diff.deleted')
ui.write('\n', label='')
ui.write('+    print "hello world"', label='diff.inserted')
ui.write('    ', label='diff.trailingwhitespace')
ui.write('\n', label='')

この部分は気をつけておく必要があるかもしれません。

文字ハイライトを行う (difflib の使い方)

さて、hg と連携してハイライトの対象のテキストを得る方法がわかったので、次は実際にハイライトをしていきましょう。
最初に考える必要があるのは類似する行を見つけることです。
たとえば次のような差分があったとき、どの行とどの行を比較して文字ハイライトをすればよいでしょうか。

- print "hello world"
- print "good-bye world"
+ print "hello WORLD"
+
+
+
+ print "good-bye WORLD"

人間が目で見ると - の 1行目と + の 1行目、- の 2行目と + の 5行目を比較すれば良いというのはすぐ判るのですが、
これを計算で求めるにはなかなかに骨が折れます。

似ている行を見つける: SequenceMatcher#ratio()

そこで登場するのが difflib です。difflib は python の標準ライブラリに含まれている差分計算用のライブラリです。
difflib には ndiff() 関数や Differ クラスなど diff とその周辺機能が揃っているので、この手の機能を作りこむ時に役に立ちます。

ここでは SequenceMatcher クラスを利用します。
SequenceMatcher はふたつのシーケンスがどのくらいマッチしているのかを知ることができます。
近似度を得るには SequenceMatcher#ratio() メソッドを利用します*1

先程の例を元に、- の 1行目と最も近似している行はどれか調べてみましょう。

>>> from difflib import SequenceMatcher
>>> matcher = SequenceMatcher()
>>> matcher.set_seq2('print "hello WORLD"')
>>> matcher.set_seq1('print "hello world"')
>>> matcher.ratio()
0.7368421052631579
>>> matcher.set_seq1('')
>>> matcher.ratio()
0.0
>>> matcher.set_seq1('print "good-bye WORLD"')
>>> matcher.ratio()
0.7317073170731707

どうやら + の 1行目、'print "hello world"' が一番近似しているようです。

この SequenceMatcher を元に、一致する行を隣り合わせにすると次のように並べ替えることができます。

- print "hello world"
+ print "hello WORLD"
+
+
+
- print "good-bye world"
+ print "good-bye WORLD"

どうでしょう。最初に予想した通りの順に並んでいませんか。

差分を見つける: SequenceMatcher#get_opcode()

SequenceMatcher#ratio() を使うことで比較対象となる行を見つけることができたので、次は文字単位のハイライトです。
SequenceMatcher の get_opcodes() メソッドはふたつの文字列の一致する箇所/異なる箇所を見つけるのにも使うことができます。

>>> from difflib import SequenceMatcher
>>> matcher = SequenceMatcher()
>>> matcher.set_seq2('print "hello WORLD"')
>>> matcher.set_seq1('print "hello world"')
>>> list(matcher.get_opcodes())
[('equal', 0, 13, 0, 13), ('replace', 13, 18, 13, 18), ('equal', 18, 19, 18, 19)]

SequenceMatcher#get_opcodes() は文字列の差分がある箇所を教えてくれます。
最初の要素 ('equal', 0, 13, 0, 13) は seq1[0:13] と seq2[0:13] の内容('print "hello ')が同じ(equal)であることを教えてくれます。
次の要素は seq1[13:18] ('world')が seq2[13:18] ('world')に置き換わっている(replace)ということを教えてくれます。
他にも tag (タプルの先頭のデータ) には 'insert' (追加された)、'deleted' (削除された) の 2種類があり、
ふたつの文字列がどのように違うのかを詳しく比較してくれます。

これらの情報を用いると文字単位でハイライトさせることができそうです。

ハイライトを調整する

というわけで、difflib.SequenceMatcher を使うことで文字単位でのハイライトを簡単に作ることができるわけですが、
そのまま素直に実装してしまうとこういう結果になります。
f:id:tk0miya:20131220005255p:plain
changes と hunk は確かに h と n が書き換わっていない(equal)なのですがこれは…

f:id:tk0miya:20131220005258p:plain
oldline と decode でこういうハイライトになるのは…
それに ('utf-8'), の閉じ括弧だけハイライトされないのも微妙な感じ…

というわけで、適当にハイライトを調整しておくと見やすくなります。
hg-diff-highlight ではヒューリスティックな関数を用意してほどよくいじっています。
f:id:tk0miya:20131220010822p:plain
f:id:tk0miya:20131220010825p:plain

まとめ

  • mercurial のカラーリングは mercurial.ui.ui クラスを通しているよ
  • difflib を使うと diff が簡単に作れる
    • 上で説明したロジックは、おおよそ Differ クラスで実装されているので参考になります。
  • difflib.SequenceMatcher は分かりづらいが、賢く動いてくれる
    • 見やすさを考えると、微調整しておくと幸せになれます

*1:リファレンスにも説明がありますが、ratio() は計算コストが高いので real_quick_ratio(), quick_ratio() を組み合わせつつ計算すると良いようです。うまい使い方は Differ#_fancy_replace() が参考になるでしょう。