diff --git a/pytools/debug.py b/pytools/debug.py index 08fbaa3354bed837763fd50a67b323d07f102a53..94490aa77150fe81def8593b9b4c6db9771df497 100644 --- a/pytools/debug.py +++ b/pytools/debug.py @@ -1,7 +1,9 @@ from __future__ import absolute_import, print_function +import sys import six from six.moves import input + from pytools import memoize @@ -216,4 +218,27 @@ def shell(locals_=None, globals_=None): # }}} + +# {{{ estimate memory usage + +def estimate_memory_usage(root, seen_ids=None): + if seen_ids is None: + seen_ids = set() + + id_root = id(root) + if id_root in seen_ids: + return 0 + + seen_ids.add(id_root) + + result = sys.getsizeof(root) + + from gc import get_referents + for ref in get_referents(root): + result += estimate_memory_usage(ref, seen_ids=seen_ids) + + return result + +# }}} + # vim: foldmethod=marker