122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
|
"""Utility for dumping memory usage stats.
|
||
|
|
||
|
This is tailored to mypy and knows (a little) about which list objects are
|
||
|
owned by particular AST nodes, etc.
|
||
|
"""
|
||
|
|
||
|
from __future__ import annotations
|
||
|
|
||
|
import gc
|
||
|
import sys
|
||
|
from collections import defaultdict
|
||
|
from typing import Dict, Iterable, cast
|
||
|
|
||
|
from mypy.nodes import FakeInfo, Node
|
||
|
from mypy.types import Type
|
||
|
from mypy.util import get_class_descriptors
|
||
|
|
||
|
|
||
|
def collect_memory_stats() -> tuple[dict[str, int], dict[str, int]]:
|
||
|
"""Return stats about memory use.
|
||
|
|
||
|
Return a tuple with these items:
|
||
|
- Dict from object kind to number of instances of that kind
|
||
|
- Dict from object kind to total bytes used by all instances of that kind
|
||
|
"""
|
||
|
objs = gc.get_objects()
|
||
|
find_recursive_objects(objs)
|
||
|
|
||
|
inferred = {}
|
||
|
for obj in objs:
|
||
|
if type(obj) is FakeInfo:
|
||
|
# Processing these would cause a crash.
|
||
|
continue
|
||
|
n = type(obj).__name__
|
||
|
if hasattr(obj, "__dict__"):
|
||
|
# Keep track of which class a particular __dict__ is associated with.
|
||
|
inferred[id(obj.__dict__)] = f"{n} (__dict__)"
|
||
|
if isinstance(obj, (Node, Type)): # type: ignore[misc]
|
||
|
if hasattr(obj, "__dict__"):
|
||
|
for x in obj.__dict__.values():
|
||
|
if isinstance(x, list):
|
||
|
# Keep track of which node a list is associated with.
|
||
|
inferred[id(x)] = f"{n} (list)"
|
||
|
if isinstance(x, tuple):
|
||
|
# Keep track of which node a list is associated with.
|
||
|
inferred[id(x)] = f"{n} (tuple)"
|
||
|
|
||
|
for k in get_class_descriptors(type(obj)):
|
||
|
x = getattr(obj, k, None)
|
||
|
if isinstance(x, list):
|
||
|
inferred[id(x)] = f"{n} (list)"
|
||
|
if isinstance(x, tuple):
|
||
|
inferred[id(x)] = f"{n} (tuple)"
|
||
|
|
||
|
freqs: dict[str, int] = {}
|
||
|
memuse: dict[str, int] = {}
|
||
|
for obj in objs:
|
||
|
if id(obj) in inferred:
|
||
|
name = inferred[id(obj)]
|
||
|
else:
|
||
|
name = type(obj).__name__
|
||
|
freqs[name] = freqs.get(name, 0) + 1
|
||
|
memuse[name] = memuse.get(name, 0) + sys.getsizeof(obj)
|
||
|
|
||
|
return freqs, memuse
|
||
|
|
||
|
|
||
|
def print_memory_profile(run_gc: bool = True) -> None:
|
||
|
if not sys.platform.startswith("win"):
|
||
|
import resource
|
||
|
|
||
|
system_memuse = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
|
||
|
else:
|
||
|
system_memuse = -1 # TODO: Support this on Windows
|
||
|
if run_gc:
|
||
|
gc.collect()
|
||
|
freqs, memuse = collect_memory_stats()
|
||
|
print("%7s %7s %7s %s" % ("Freq", "Size(k)", "AvgSize", "Type"))
|
||
|
print("-------------------------------------------")
|
||
|
totalmem = 0
|
||
|
i = 0
|
||
|
for n, mem in sorted(memuse.items(), key=lambda x: -x[1]):
|
||
|
f = freqs[n]
|
||
|
if i < 50:
|
||
|
print("%7d %7d %7.0f %s" % (f, mem // 1024, mem / f, n))
|
||
|
i += 1
|
||
|
totalmem += mem
|
||
|
print()
|
||
|
print("Mem usage RSS ", system_memuse // 1024)
|
||
|
print("Total reachable ", totalmem // 1024)
|
||
|
|
||
|
|
||
|
def find_recursive_objects(objs: list[object]) -> None:
|
||
|
"""Find additional objects referenced by objs and append them to objs.
|
||
|
|
||
|
We use this since gc.get_objects() does not return objects without pointers
|
||
|
in them such as strings.
|
||
|
"""
|
||
|
seen = {id(o) for o in objs}
|
||
|
|
||
|
def visit(o: object) -> None:
|
||
|
if id(o) not in seen:
|
||
|
objs.append(o)
|
||
|
seen.add(id(o))
|
||
|
|
||
|
for obj in objs.copy():
|
||
|
if type(obj) is FakeInfo:
|
||
|
# Processing these would cause a crash.
|
||
|
continue
|
||
|
if type(obj) in (dict, defaultdict):
|
||
|
for key, val in cast(Dict[object, object], obj).items():
|
||
|
visit(key)
|
||
|
visit(val)
|
||
|
if type(obj) in (list, tuple, set):
|
||
|
for x in cast(Iterable[object], obj):
|
||
|
visit(x)
|
||
|
if hasattr(obj, "__slots__"):
|
||
|
for base in type.mro(type(obj)):
|
||
|
for slot in getattr(base, "__slots__", ()):
|
||
|
if hasattr(obj, slot):
|
||
|
visit(getattr(obj, slot))
|