Class PstatsCollector

Lib/profiling/sampling/pstats_collector.py:10–422

class PstatsCollector(Collector):
    def __init__(self, sample_interval_usec, *, skip_idle=False):
        self.result = collections.defaultdict(
            lambda: dict(total_rec_calls=0, direct_calls=0, cumulative_calls=0)
        )
        self.stats = {}
        self.sample_interval_usec = sample_interval_usec
        self.callers = collections.defaultdict(
            lambda: collections.defaultdict(int)
        )
        self.skip_idle = skip_idle
        self._seen_locations = set()

    def _process_frames(self, frames, weight=1):
        """Process a single thread's frame stack."""
        if not frames:
            return

        self._seen_locations.clear()

        # Process each frame in the stack to track cumulative calls
        # frame.location is int, tuple (lineno, end_lineno, col_offset, end_col_offset), or None
        for frame in frames:
            lineno = extract_lineno(frame.location)
            location = (frame.filename, lineno, frame.funcname)
            if location not in self._seen_locations:
                self._seen_locations.add(location)
                self.result[location]["cumulative_calls"] += weight

        # The top frame gets counted as an inline call (directly executing)
        top_lineno = extract_lineno(frames[0].location)
        top_location = (frames[0].filename, top_lineno, frames[0].funcname)
        self.result[top_location]["direct_calls"] += weight

        # Track caller-callee relationships for call graph
        for i in range(1, len(frames)):
            callee_frame = frames[i - 1]
            caller_frame = frames[i]

            callee_lineno = extract_lineno(callee_frame.location)
            caller_lineno = extract_lineno(caller_frame.location)
            callee = (callee_frame.filename, callee_lineno, callee_frame.funcname)
            caller = (caller_frame.filename, caller_lineno, caller_frame.funcname)

            self.callers[callee][caller] += weight

    def collect(self, stack_frames, timestamps_us=None):
        weight = len(timestamps_us) if timestamps_us else 1
        for frames, _ in self._iter_stacks(stack_frames, skip_idle=self.skip_idle):
            self._process_frames(frames, weight=weight)

    def export(self, filename):
        self.create_stats()
        self._dump_stats(filename)

    def _dump_stats(self, file):
        stats_with_marker = dict(self.stats)
        stats_with_marker[("__sampled__",)] = True