現在Pythonのプログラム上で、どの変数がどれくらいのメモリを確保しているのかを確認するためのモジュールとして、tracemallocという標準モジュールがあります。

これを使えば例えばこんなコードを書くことで、ある処理のメモリ解放漏れを確認することができます。

tracemallocを直接呼び出すコード

import tracemalloc

tracemalloc.start()
snap1 = tracemalloc.take_snapshot()
size1 = sum([stat.size for stat in snap1.statistics("filename")])
for stat in snap1.statistics('lineno')[:10]:
  print(stat)
# ... ある処理 ...
snap2 = tracemalloc.take_snapshot()
size2 = sum([stat.size for stat in snap2.statistics("filename")])
for stat in snap2.compare_to(snap1, 'lineno')[:10]:
  print(stat)
diff = abs(size1 - size2)
print("処理終了後メモリ量{:,}バイト".format(diff))

with構文を使うなら、こんな感じ

with構文を使うべくクラス化したもの

import tracemalloc

class MemCheck:
  """
  メモリチェックを行うためのクラス。with構文で使用する。
  withを抜けたときに自動的にチェックが行われる。
  """

  def __init__(self):
    """
    初期化
    """
    pass

  def __enter__(self):
    tracemalloc.start()
    self.snap = tracemalloc.take_snapshot()
    self.size = sum([stat.size for stat in self.snap.statistics("filename")])
    print("")
    print("-----TEST START!!!!-----")
    for stat in self.snap.statistics('lineno')[:10]:
      print(stat)
    return self

  def __exit__(self, ex_type, ex_value, trace):
    print("-----TEST END!!!!!!-----")
    snap = tracemalloc.take_snapshot()
    size = sum([stat.size for stat in snap.statistics("filename")])
    for stat in snap.compare_to(self.snap, 'lineno')[:10]:
      print(stat)
    diff = abs(self.size - size)
    print("処理終了後メモリ量{:,}バイト".format(diff))
    print("-" * 20)
    return False

def main():
  with MemCheck():
    # ...ある処理...

ただ実際これを動かすと、メモリを使用しているオブジェクトの中に、tracemallocモジュール内部や、デバッガの動作用のオブジェクトも含まれてしまうため、メモリ解放漏れのテストなどに使おうとするとやや厄介です。

そのために、データをフィルタするためのfilter_traces()というメソッドが用意されています。

tracemalloc --- メモリ割り当ての追跡 — Python 3.8.3 ドキュメント

ですが、ドキュメントを見ると用例が二つしか書かれていないため、「デバッグ用のオブジェクトはまるっと無視したい」とか、逆に「自分が書いたコードだけをテスト対象にしたい」という事態に対応できません。

フォルダ名を指定して「このフォルダ配下のファイルだけをトレースしたい」とかいうときにはどうすればいいのか。

こうする

以下のように書けばいいのです。

with構文を使う例にフィルタを追加したもの

import tracemalloc
from pathlib import Path

class MemCheck:
  """
  メモリチェックを行うためのクラス。with構文で使用する。
  withを抜けたときに自動的にチェックが行われる。
  """

  def __init__(self):
    """
    初期化
    """
    pass

  def __enter__(self):
    tracemalloc.start()
    self.snap = tracemalloc.take_snapshot().filter_traces(self.get_filter_traces())
    self.size = sum([stat.size for stat in self.snap.statistics("filename")])
    print("")
    print("-----TEST START!!!!-----")
    for stat in self.snap.statistics('lineno')[:10]:
      print(stat)
    return self

  def __exit__(self, ex_type, ex_value, trace):
    print("-----TEST END!!!!!!-----")
    snap = tracemalloc.take_snapshot().filter_traces(self.get_filter_traces())
    size = sum([stat.size for stat in snap.statistics("filename")])
    for stat in snap.compare_to(self.snap, 'lineno')[:10]:
      print(stat)
    diff = abs(self.size - size)
    print("処理終了後メモリ量{:,}バイト".format(diff))
    print("-" * 20)
    return False

  # 追加
  def get_filter_traces(self):
    return (
      tracemalloc.Filter(True, str(Path(__file__).parent.parent / ".venv" / "lib" / "site-packages" / "*")),
      tracemalloc.Filter(True, str(Path(__file__).parent.parent / "src" / "*")),
    )

def main():
  with MemCheck():
    # ...ある処理...

コード内にあるget_filter_traces()メソッドが、filter_traces()メソッド用のフィルタリストを返却するメソッドです。フィルタは全てのスナップショット取得処理に適用させたいので、このようにしています。

filter_traces()メソッドの引数はtracemalloc.Filterオブジェクトのタプルを受け取るようになっており、tracemalloc.Filterオブジェクトは、コンストラクタの第一引数に、そのフィルタに「一致するものを表示する(True)」のか、「一致しないものを表示する(False)」を指定、第二引数(filename_pattern)にフィルタに指定するファイル名を指定します。

このファイル名、ドキュメントを見ると「Filename pattern of the filter (str). Read-only property.」と書いてあるので一見正規表現でも突っ込めばいいのかというと、そうではなく、fnmatchというモジュールで処理できるShell形式のワイルドカード文字列を指定します)

正規表現などのパターンマッチを使い慣れてるとパターン＝正規表現と思ってしまう人もいるかもしれませんが、正規表現ではないのでご注意ください。であればそれくらい分かるようなコメントを書いて欲しかったところですが・・・。

srcフォルダ配下のファイルだけをtrace対象にすると何も出てこなくなることがある

なお、メモリを確保したのが自分が書いたプログラムそのものではなく自分が書いたプログラムで呼び出したオブジェクトである場合は、tracemalloc.take_snapshot()で出てくるのは自分で書いたソースファイルの行番号ではなく、呼び出したオブジェクトを定義しているソースファイルの行番号になります(自分のプログラムでPyPDF.PdfFileReaderを呼び出し、そこで読み込んだメモリが解放されていなかった場合、「\lib\site-packages\PyPDF2\配下のファイルでメモリが確保されている」と表示される)。

このため、フィルタを作成する場合は「srcフォルダ配下」だけでなく、上記例のように「呼び出しているモジュールのフォルダ」もtrace対象に入れるようにしましょう。仮想環境のフォルダがプロジェクト内に生成されるようになっていれば、上記のように.venv\lib\site-packages\*あたりが検索対象になっていればOKです。

どうやって気付いたのか

C:\PythonNN\Lib\tracemalloc.pyファイルを開いてFilterの処理を追いました(棒)。

こういうことしなければ分からない挙動は勘弁して欲しい…。