RustPython/Lib/profiling/sampling/pstats_collector.py

import collections
import marshal
import pstats
lazy from _colorize import get_colors

from .collector import Collector, extract_lineno
from .constants import MICROSECONDS_PER_SECOND, PROFILING_MODE_CPU


class PstatsCollector(Collector):
    aggregating = True

    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)
        return True

    def _dump_stats(self, file):
        stats_with_marker = dict(self.stats)
        stats_with_marker[("__sampled__",)] = True
        with open(file, "wb") as f:
            marshal.dump(stats_with_marker, f)

    # Needed for compatibility with pstats.Stats
    def create_stats(self):
        sample_interval_sec = self.sample_interval_usec / MICROSECONDS_PER_SECOND
        callers = {}
        for fname, call_counts in self.result.items():
            total = call_counts["direct_calls"] * sample_interval_sec
            cumulative_calls = call_counts["cumulative_calls"]
            cumulative = cumulative_calls * sample_interval_sec
            callers = dict(self.callers.get(fname, {}))
            self.stats[fname] = (
                call_counts["direct_calls"],  # cc = direct calls for sample percentage
                cumulative_calls,  # nc = cumulative calls for cumulative percentage
                total,
                cumulative,
                callers,
            )

    def print_stats(self, sort=-1, limit=None, show_summary=True, mode=None):
        """Print formatted statistics to stdout."""
        # Create stats object
        stats = pstats.SampledStats(self).strip_dirs()
        if not stats.stats:
            print("No samples were collected.")
            if mode == PROFILING_MODE_CPU:
                print("This can happen in CPU mode when all threads are idle.")
            return

        # Get the stats data
        stats_list = []
        for func, (
            direct_calls,
            cumulative_calls,
            total_time,
            cumulative_time,
            callers,
        ) in stats.stats.items():
            stats_list.append(
                (
                    func,
                    direct_calls,
                    cumulative_calls,
                    total_time,
                    cumulative_time,
                    callers,
                )
            )

        # Calculate total samples for percentage calculations (using direct_calls)
        total_samples = sum(
            direct_calls for _, direct_calls, _, _, _, _ in stats_list
        )

        # Sort based on the requested field
        sort_field = sort
        if sort_field == -1:  # stdname
            stats_list.sort(key=lambda x: str(x[0]))
        elif sort_field == 0:  # nsamples (direct samples)
            stats_list.sort(key=lambda x: x[1], reverse=True)  # direct_calls
        elif sort_field == 1:  # tottime
            stats_list.sort(key=lambda x: x[3], reverse=True)  # total_time
        elif sort_field == 2:  # cumtime
            stats_list.sort(key=lambda x: x[4], reverse=True)  # cumulative_time
        elif sort_field == 3:  # sample%
            stats_list.sort(
                key=lambda x: (x[1] / total_samples * 100)
                if total_samples > 0
                else 0,
                reverse=True,  # direct_calls percentage
            )
        elif sort_field == 4:  # cumul%
            stats_list.sort(
                key=lambda x: (x[2] / total_samples * 100)
                if total_samples > 0
                else 0,
                reverse=True,  # cumulative_calls percentage
            )
        elif sort_field == 5:  # nsamples (cumulative samples)
            stats_list.sort(key=lambda x: x[2], reverse=True)  # cumulative_calls

        # Apply limit if specified
        if limit is not None:
            stats_list = stats_list[:limit]

        # Determine the best unit for time columns based on maximum values
        max_total_time = max(
            (total_time for _, _, _, total_time, _, _ in stats_list), default=0
        )
        max_cumulative_time = max(
            (cumulative_time for _, _, _, _, cumulative_time, _ in stats_list),
            default=0,
        )

        total_time_unit, total_time_scale = self._determine_best_unit(max_total_time)
        cumulative_time_unit, cumulative_time_scale = self._determine_best_unit(
            max_cumulative_time
        )

        # Define column widths for consistent alignment
        col_widths = {
            "nsamples": 15,  # "nsamples" column (inline/cumulative format)
            "sample_pct": 8,  # "sample%" column
            "tottime": max(12, len(f"tottime ({total_time_unit})")),
            "cum_pct": 8,  # "cumul%" column
            "cumtime": max(12, len(f"cumtime ({cumulative_time_unit})")),
        }

        # Print header with colors and proper alignment
        ANSIColors = get_colors()

        print(f"{ANSIColors.BOLD_BLUE}Profile Stats:{ANSIColors.RESET}")

        header_nsamples = f"{ANSIColors.BOLD_BLUE}{'nsamples':>{col_widths['nsamples']}}{ANSIColors.RESET}"
        header_sample_pct = f"{ANSIColors.BOLD_BLUE}{'sample%':>{col_widths['sample_pct']}}{ANSIColors.RESET}"
        header_tottime = f"{ANSIColors.BOLD_BLUE}{f'tottime ({total_time_unit})':>{col_widths['tottime']}}{ANSIColors.RESET}"
        header_cum_pct = f"{ANSIColors.BOLD_BLUE}{'cumul%':>{col_widths['cum_pct']}}{ANSIColors.RESET}"
        header_cumtime = f"{ANSIColors.BOLD_BLUE}{f'cumtime ({cumulative_time_unit})':>{col_widths['cumtime']}}{ANSIColors.RESET}"
        header_filename = (
            f"{ANSIColors.BOLD_BLUE}filename:lineno(function){ANSIColors.RESET}"
        )

        print(
            f"{header_nsamples}  {header_sample_pct}  {header_tottime}  {header_cum_pct}  {header_cumtime}  {header_filename}"
        )

        # Print each line with proper alignment
        for (
            func,
            direct_calls,
            cumulative_calls,
            total_time,
            cumulative_time,
            callers,
        ) in stats_list:
            # Calculate percentages
            sample_pct = (
                (direct_calls / total_samples * 100) if total_samples > 0 else 0
            )
            cum_pct = (
                (cumulative_calls / total_samples * 100)
                if total_samples > 0
                else 0
            )

            # Format values with proper alignment - always use A/B format
            nsamples_str = f"{direct_calls}/{cumulative_calls}"
            nsamples_str = f"{nsamples_str:>{col_widths['nsamples']}}"
            sample_pct_str = f"{sample_pct:{col_widths['sample_pct']}.1f}"
            tottime = f"{total_time * total_time_scale:{col_widths['tottime']}.3f}"
            cum_pct_str = f"{cum_pct:{col_widths['cum_pct']}.1f}"
            cumtime = f"{cumulative_time * cumulative_time_scale:{col_widths['cumtime']}.3f}"

            # Format the function name with colors
            func_name = (
                f"{ANSIColors.GREEN}{func[0]}{ANSIColors.RESET}:"
                f"{ANSIColors.YELLOW}{func[1]}{ANSIColors.RESET}("
                f"{ANSIColors.CYAN}{func[2]}{ANSIColors.RESET})"
            )

            # Print the formatted line with consistent spacing
            print(
                f"{nsamples_str}  {sample_pct_str}  {tottime}  {cum_pct_str}  {cumtime}  {func_name}"
            )

        # Print legend
        print(f"\n{ANSIColors.BOLD_BLUE}Legend:{ANSIColors.RESET}")
        print(
            f"  {ANSIColors.YELLOW}nsamples{ANSIColors.RESET}: Direct/Cumulative samples (direct executing / on call stack)"
        )
        print(
            f"  {ANSIColors.YELLOW}sample%{ANSIColors.RESET}: Percentage of total samples this function was directly executing"
        )
        print(
            f"  {ANSIColors.YELLOW}tottime{ANSIColors.RESET}: Estimated total time spent directly in this function"
        )
        print(
            f"  {ANSIColors.YELLOW}cumul%{ANSIColors.RESET}: Percentage of total samples when this function was on the call stack"
        )
        print(
            f"  {ANSIColors.YELLOW}cumtime{ANSIColors.RESET}: Estimated cumulative time (including time in called functions)"
        )
        print(
            f"  {ANSIColors.YELLOW}filename:lineno(function){ANSIColors.RESET}: Function location and name"
        )

        # Print summary of interesting functions if enabled
        if show_summary and stats_list:
            self._print_summary(stats_list, total_samples)

    @staticmethod
    def _determine_best_unit(max_value):
        """Determine the best unit (s, ms, μs) and scale factor for a maximum value."""
        if max_value >= 1.0:
            return "s", 1.0
        elif max_value >= 0.001:
            return "ms", 1000.0
        else:
            return "μs", float(MICROSECONDS_PER_SECOND)

    def _print_summary(self, stats_list, total_samples):
        """Print summary of interesting functions."""
        ANSIColors = get_colors()

        print(
            f"\n{ANSIColors.BOLD_BLUE}Summary of Interesting Functions:{ANSIColors.RESET}"
        )

        # Aggregate stats by fully qualified function name (ignoring line numbers)
        func_aggregated = {}
        for (
            func,
            direct_calls,
            cumulative_calls,
            total_time,
            cumulative_time,
            callers,
        ) in stats_list:
            # Use filename:function_name as the key to get fully qualified name
            qualified_name = f"{func[0]}:{func[2]}"
            if qualified_name not in func_aggregated:
                func_aggregated[qualified_name] = [
                    0,
                    0,
                    0,
                    0,
                ]  # direct_calls, cumulative_calls, total_time, cumulative_time
            func_aggregated[qualified_name][0] += direct_calls
            func_aggregated[qualified_name][1] += cumulative_calls
            func_aggregated[qualified_name][2] += total_time
            func_aggregated[qualified_name][3] += cumulative_time

        # Convert aggregated data back to list format for processing
        aggregated_stats = []
        for qualified_name, (
            prim_calls,
            total_calls,
            total_time,
            cumulative_time,
        ) in func_aggregated.items():
            # Parse the qualified name back to filename and function name
            if ":" in qualified_name:
                filename, func_name = qualified_name.rsplit(":", 1)
            else:
                filename, func_name = "", qualified_name
            # Create a dummy func tuple with filename and function name for display
            dummy_func = (filename, "", func_name)
            aggregated_stats.append(
                (
                    dummy_func,
                    prim_calls,
                    total_calls,
                    total_time,
                    cumulative_time,
                    {},
                )
            )

        # Determine best units for summary metrics
        max_total_time = max(
            (total_time for _, _, _, total_time, _, _ in aggregated_stats),
            default=0,
        )
        max_cumulative_time = max(
            (
                cumulative_time
                for _, _, _, _, cumulative_time, _ in aggregated_stats
            ),
            default=0,
        )

        total_unit, total_scale = self._determine_best_unit(max_total_time)
        cumulative_unit, cumulative_scale = self._determine_best_unit(
            max_cumulative_time
        )

        def _format_func_name(func):
            """Format function name with colors."""
            return (
                f"{ANSIColors.GREEN}{func[0]}{ANSIColors.RESET}:"
                f"{ANSIColors.YELLOW}{func[1]}{ANSIColors.RESET}("
                f"{ANSIColors.CYAN}{func[2]}{ANSIColors.RESET})"
            )

        def _print_top_functions(stats_list, title, key_func, format_line, n=3):
            """Print top N functions sorted by key_func with formatted output."""
            print(f"\n{ANSIColors.BOLD_BLUE}{title}:{ANSIColors.RESET}")
            sorted_stats = sorted(stats_list, key=key_func, reverse=True)
            for stat in sorted_stats[:n]:
                if line := format_line(stat):
                    print(f"  {line}")

        # Functions with highest direct/cumulative ratio (hot spots)
        def format_hotspots(stat):
            func, direct_calls, cumulative_calls, total_time, _, _ = stat
            if direct_calls > 0 and cumulative_calls > 0:
                ratio = direct_calls / cumulative_calls
                direct_pct = (
                    (direct_calls / total_samples * 100)
                    if total_samples > 0
                    else 0
                )
                return (
                    f"{ratio:.3f} direct/cumulative ratio, "
                    f"{direct_pct:.1f}% direct samples: {_format_func_name(func)}"
                )
            return None

        _print_top_functions(
            aggregated_stats,
            "Functions with Highest Direct/Cumulative Ratio (Hot Spots)",
            key_func=lambda x: (x[1] / x[2]) if x[2] > 0 else 0,
            format_line=format_hotspots,
        )

        # Functions with highest call frequency (cumulative/direct difference)
        def format_call_frequency(stat):
            func, direct_calls, cumulative_calls, total_time, _, _ = stat
            if cumulative_calls > direct_calls:
                call_frequency = cumulative_calls - direct_calls
                cum_pct = (
                    (cumulative_calls / total_samples * 100)
                    if total_samples > 0
                    else 0
                )
                return (
                    f"{call_frequency:d} indirect calls, "
                    f"{cum_pct:.1f}% total stack presence: {_format_func_name(func)}"
                )
            return None

        _print_top_functions(
            aggregated_stats,
            "Functions with Highest Call Frequency (Indirect Calls)",
            key_func=lambda x: x[2] - x[1],  # Sort by (cumulative - direct)
            format_line=format_call_frequency,
        )

        # Functions with highest cumulative-to-direct multiplier (call magnification)
        def format_call_magnification(stat):
            func, direct_calls, cumulative_calls, total_time, _, _ = stat
            if direct_calls > 0 and cumulative_calls > direct_calls:
                multiplier = cumulative_calls / direct_calls
                indirect_calls = cumulative_calls - direct_calls
                return (
                    f"{multiplier:.1f}x call magnification, "
                    f"{indirect_calls:d} indirect calls from {direct_calls:d} direct: {_format_func_name(func)}"
                )
            return None

        _print_top_functions(
            aggregated_stats,
            "Functions with Highest Call Magnification (Cumulative/Direct)",
            key_func=lambda x: (x[2] / x[1])
            if x[1] > 0
            else 0,  # Sort by cumulative/direct ratio
            format_line=format_call_magnification,
        )