from collections.abc import Sequence, Iterable from functools import total_ordering import fnmatch import linecache import os.path import pickle # Import types and functions implemented in C from _tracemalloc import * from _tracemalloc import _get_object_traceback, _get_traces def _format_size(size, sign): for unit in ('B', 'KiB', 'MiB', 'GiB', 'TiB'): if abs(size) < 100 and unit != 'B': # 3 digits (xx.x UNIT) if sign: return "%+.1f %s" % (size, unit) else: return "%.1f %s" % (size, unit) if abs(size) < 10 * 1024 or unit == 'TiB': # 4 or 5 digits (xxxx UNIT) if sign: return "%+.0f %s" % (size, unit) else: return "%.0f %s" % (size, unit) size /= 1024 class Statistic: """ Statistic difference on memory allocations between two Snapshot instance. """ __slots__ = ('traceback', 'size', 'count') def __init__(self, traceback, size, count): self.traceback = traceback self.size = size self.count = count def __hash__(self): return hash((self.traceback, self.size, self.count)) def __eq__(self, other): if not isinstance(other, Statistic): return NotImplemented return (self.traceback == other.traceback and self.size == other.size and self.count == other.count) def __str__(self): text = ("%s: size=%s, count=%i" % (self.traceback, _format_size(self.size, False), self.count)) if self.count: average = self.size / self.count text += ", average=%s" % _format_size(average, False) return text def __repr__(self): return ('<Statistic traceback=%r size=%i count=%i>' % (self.traceback, self.size, self.count)) def _sort_key(self): return (self.size, self.count, self.traceback) class StatisticDiff: """ Statistic difference on memory allocations between an old and a new Snapshot instance. """ __slots__ = ('traceback', 'size', 'size_diff', 'count', 'count_diff') def __init__(self, traceback, size, size_diff, count, count_diff): self.traceback = traceback self.size = size self.size_diff = size_diff self.count = count self.count_diff = count_diff def __hash__(self): return hash((self.traceback, self.size, self.size_diff, self.count, self.count_diff)) def __eq__(self, other): if not isinstance(other, StatisticDiff): return NotImplemented return (self.traceback == other.traceback and self.size == other.size and self.size_diff == other.size_diff and self.count == other.count and self.count_diff == other.count_diff) def __str__(self): text = ("%s: size=%s (%s), count=%i (%+i)" % (self.traceback, _format_size(self.size, False), _format_size(self.size_diff, True), self.count, self.count_diff)) if self.count: average = self.size / self.count text += ", average=%s" % _format_size(average, False) return text def __repr__(self): return ('<StatisticDiff traceback=%r size=%i (%+i) count=%i (%+i)>' % (self.traceback, self.size, self.size_diff, self.count, self.count_diff)) def _sort_key(self): return (abs(self.size_diff), self.size, abs(self.count_diff), self.count, self.traceback) def _compare_grouped_stats(old_group, new_group): statistics = [] for traceback, stat in new_group.items(): previous = old_group.pop(traceback, None) if previous is not None: stat = StatisticDiff(traceback, stat.size, stat.size - previous.size, stat.count, stat.count - previous.count) else: stat = StatisticDiff(traceback, stat.size, stat.size, stat.count, stat.count) statistics.append(stat) for traceback, stat in old_group.items(): stat = StatisticDiff(traceback, 0, -stat.size, 0, -stat.count) statistics.append(stat) return statistics @total_ordering class Frame: """ Frame of a traceback. """ __slots__ = ("_frame",) def __init__(self, frame): # frame is a tuple: (filename: str, lineno: int) self._frame = frame @property def filename(self): return self._frame[0] @property def lineno(self): return self._frame[1] def __eq__(self, other): if not isinstance(other, Frame): return NotImplemented return (self._frame == other._frame) def __lt__(self, other): if not isinstance(other, Frame): return NotImplemented return (self._frame < other._frame) def __hash__(self): return hash(self._frame) def __str__(self): return "%s:%s" % (self.filename, self.lineno) def __repr__(self): return "<Frame filename=%r lineno=%r>" % (self.filename, self.lineno) @total_ordering class Traceback(Sequence): """ Sequence of Frame instances sorted from the oldest frame to the most recent frame. """ __slots__ = ("_frames", '_total_nframe') def __init__(self, frames, total_nframe=None): Sequence.__init__(self) # frames is a tuple of frame tuples: see Frame constructor for the # format of a frame tuple; it is reversed, because _tracemalloc # returns frames sorted from most recent to oldest, but the # Python API expects oldest to most recent self._frames = tuple(reversed(frames)) self._total_nframe = total_nframe @property def total_nframe(self): return self._total_nframe def __len__(self): return len(self._frames) def __getitem__(self, index): if isinstance(index, slice): return tuple(Frame(trace) for trace in self._frames[index]) else: return Frame(self._frames[index]) def __contains__(self, frame): return frame._frame in self._frames def __hash__(self): return hash(self._frames) def __eq__(self, other): if not isinstance(other, Traceback): return NotImplemented return (self._frames == other._frames) def __lt__(self, other): if not isinstance(other, Traceback): return NotImplemented return (self._frames < other._frames) def __str__(self): return str(self[0]) def __repr__(self): s = f"<Traceback {tuple(self)}" if self._total_nframe is None: s += ">" else: s += f" total_nframe={self.total_nframe}>" return s def format(self, limit=None, most_recent_first=False): lines = [] if limit is not None: if limit > 0: frame_slice = self[-limit:] else: frame_slice = self[:limit] else: frame_slice = self if most_recent_first: frame_slice = reversed(frame_slice) for frame in frame_slice: lines.append(' File "%s", line %s' % (frame.filename, frame.lineno)) line = linecache.getline(frame.filename, frame.lineno).strip() if line: lines.append(' %s' % line) return lines def get_object_traceback(obj): """ Get the traceback where the Python object *obj* was allocated. Return a Traceback instance. Return None if the tracemalloc module is not tracing memory allocations or did not trace the allocation of the object. """ frames = _get_object_traceback(obj) if frames is not None: return Traceback(frames) else: return None class Trace: """ Trace of a memory block. """ __slots__ = ("_trace",) def __init__(self, trace): # trace is a tuple: (domain: int, size: int, traceback: tuple). # See Traceback constructor for the format of the traceback tuple. self._trace = trace @property def domain(self): return self._trace[0] @property def size(self): return self._trace[1] @property def traceback(self): return Traceback(*self._trace[2:]) def __eq__(self, other): if not isinstance(other, Trace): return NotImplemented return (self._trace == other._trace) def __hash__(self): return hash(self._trace) def __str__(self): return "%s: %s" % (self.traceback, _format_size(self.size, False)) def __repr__(self): return ("<Trace domain=%s size=%s, traceback=%r>" % (self.domain, _format_size(self.size, False), self.traceback)) class _Traces(Sequence): def __init__(self, traces): Sequence.__init__(self) # traces is a tuple of trace tuples: see Trace constructor self._traces = traces def __len__(self): return len(self._traces) def __getitem__(self, index): if isinstance(index, slice): return tuple(Trace(trace) for trace in self._traces[index]) else: return Trace(self._traces[index]) def __contains__(self, trace): return trace._trace in self._traces def __eq__(self, other): if not isinstance(other, _Traces): return NotImplemented return (self._traces == other._traces) def __repr__(self): return "<Traces len=%s>" % len(self) def _normalize_filename(filename): filename = os.path.normcase(filename) if filename.endswith('.pyc'): filename = filename[:-1] return filename class BaseFilter: def __init__(self, inclusive): self.inclusive = inclusive def _match(self, trace): raise NotImplementedError class Filter(BaseFilter): def __init__(self, inclusive, filename_pattern, lineno=None, all_frames=False, domain=None): super().__init__(inclusive) self.inclusive = inclusive self._filename_pattern = _normalize_filename(filename_pattern) self.lineno = lineno self.all_frames = all_frames self.domain = domain @property def filename_pattern(self): return self._filename_pattern def _match_frame_impl(self, filename, lineno): filename = _normalize_filename(filename) if not fnmatch.fnmatch(filename, self._filename_pattern): return False if self.lineno is None: return True else: return (lineno == self.lineno) def _match_frame(self, filename, lineno): return self._match_frame_impl(filename, lineno) ^ (not self.inclusive) def _match_traceback(self, traceback): if self.all_frames: if any(self._match_frame_impl(filename, lineno) for filename, lineno in traceback): return self.inclusive else: return (not self.inclusive) else: filename, lineno = traceback[0] return self._match_frame(filename, lineno) def _match(self, trace): domain, size, traceback, total_nframe = trace res = self._match_traceback(traceback) if self.domain is not None: if self.inclusive: return res and (domain == self.domain) else: return res or (domain != self.domain) return res class DomainFilter(BaseFilter): def __init__(self, inclusive, domain): super().__init__(inclusive) self._domain = domain @property def domain(self): return self._domain def _match(self, trace): domain, size, traceback, total_nframe = trace return (domain == self.domain) ^ (not self.inclusive) class Snapshot: """ Snapshot of traces of memory blocks allocated by Python. """ def __init__(self, traces, traceback_limit): # traces is a tuple of trace tuples: see _Traces constructor for # the exact format self.traces = _Traces(traces) self.traceback_limit = traceback_limit def dump(self, filename): """ Write the snapshot into a file. """ with open(filename, "wb") as fp: pickle.dump(self, fp, pickle.HIGHEST_PROTOCOL) @staticmethod def load(filename): """ Load a snapshot from a file. """ with open(filename, "rb") as fp: return pickle.load(fp) def _filter_trace(self, include_filters, exclude_filters, trace): if include_filters: if not any(trace_filter._match(trace) for trace_filter in include_filters): return False if exclude_filters: if any(not trace_filter._match(trace) for trace_filter in exclude_filters): return False return True def filter_traces(self, filters): """ Create a new Snapshot instance with a filtered traces sequence, filters is a list of Filter or DomainFilter instances. If filters is an empty list, return a new Snapshot instance with a copy of the traces. """ if not isinstance(filters, Iterable): raise TypeError("filters must be a list of filters, not %s" % type(filters).__name__) if filters: include_filters = [] exclude_filters = [] for trace_filter in filters: if trace_filter.inclusive: include_filters.append(trace_filter) else: exclude_filters.append(trace_filter) new_traces = [trace for trace in self.traces._traces if self._filter_trace(include_filters, exclude_filters, trace)] else: new_traces = self.traces._traces.copy() return Snapshot(new_traces, self.traceback_limit) def _group_by(self, key_type, cumulative): if key_type not in ('traceback', 'filename', 'lineno'): raise ValueError("unknown key_type: %r" % (key_type,)) if cumulative and key_type not in ('lineno', 'filename'): raise ValueError("cumulative mode cannot by used " "with key type %r" % key_type) stats = {} tracebacks = {} if not cumulative: for trace in self.traces._traces: domain, size, trace_traceback, total_nframe = trace try: traceback = tracebacks[trace_traceback] except KeyError: if key_type == 'traceback': frames = trace_traceback elif key_type == 'lineno': frames = trace_traceback[:1] else: # key_type == 'filename': frames = ((trace_traceback[0][0], 0),) traceback = Traceback(frames) tracebacks[trace_traceback] = traceback try: stat = stats[traceback] stat.size += size stat.count += 1 except KeyError: stats[traceback] = Statistic(traceback, size, 1) else: # cumulative statistics for trace in self.traces._traces: domain, size, trace_traceback, total_nframe = trace for frame in trace_traceback: try: traceback = tracebacks[frame] except KeyError: if key_type == 'lineno': frames = (frame,) else: # key_type == 'filename': frames = ((frame[0], 0),) traceback = Traceback(frames) tracebacks[frame] = traceback try: stat = stats[traceback] stat.size += size stat.count += 1 except KeyError: stats[traceback] = Statistic(traceback, size, 1) return stats def statistics(self, key_type, cumulative=False): """ Group statistics by key_type. Return a sorted list of Statistic instances. """ grouped = self._group_by(key_type, cumulative) statistics = list(grouped.values()) statistics.sort(reverse=True, key=Statistic._sort_key) return statistics def compare_to(self, old_snapshot, key_type, cumulative=False): """ Compute the differences with an old snapshot old_snapshot. Get statistics as a sorted list of StatisticDiff instances, grouped by group_by. """ new_group = self._group_by(key_type, cumulative) old_group = old_snapshot._group_by(key_type, cumulative) statistics = _compare_grouped_stats(old_group, new_group) statistics.sort(reverse=True, key=StatisticDiff._sort_key) return statistics def take_snapshot(): """ Take a snapshot of traces of memory blocks allocated by Python. """ if not is_tracing(): raise RuntimeError("the tracemalloc module must be tracing memory " "allocations to take a snapshot") traces = _get_traces() traceback_limit = get_traceback_limit() return Snapshot(traces, traceback_limit)