gh-125855: pstats.Stats: Allow sorting/filtering by caller/callee

Doc/library/profile.rst

@@ -196,7 +196,13 @@ is still sorted according to the last criteria) do::
 
    p.print_callers(.5, 'init')
 
-and you would get a list of callers for each of the listed functions.
+and you would get a list of all the callers for each of the listed functions.
+That may be a very long list, so if you want to sort or limit each function's
+list of callers, you might put::
+
+   p.print_callers(.5, 'init', callers_sort=SortKey.TIME, callers_filter=5)
+
+to get the top 5 callers by time of each ``init`` function in the top 50%.
 
 If you want more functionality, you're going to have to read the manual, or
 guess what the following functions do::
@@ -507,7 +513,7 @@ Analysis of the profiler data is done using the :class:`~pstats.Stats` class.
       and then proceed to only print the first 10% of them.
 
 
-   .. method:: print_callers(*restrictions)
+   .. method:: print_callers(*restrictions, callers_sort_key=(), callers_filter=())
 
       This method for the :class:`Stats` class prints a list of all functions
       that called each function in the profiled database.  The ordering is
@@ -526,13 +532,22 @@ Analysis of the profiler data is done using the :class:`~pstats.Stats` class.
         cumulative times spent in the current function while it was invoked by
         this specific caller.
 
+      By default, all callers of a function are printed in lexicographic order.
+      To sort them differently, you can supply sort criteria in the keyword
+      argument ``callers_sort_key``. These use exactly the same format as in
+      :meth:`~pstats.Stats.sort_stats`, though you must supply them each time
+      you call this method. And to limit the number of callers printed per
+      function, you can supply restrictions in the keyword argument
+      ``callers_filter``, as you would in :meth:`~pstats.Stats.print_stats`.
+
 
-   .. method:: print_callees(*restrictions)
+   .. method:: print_callees(*restrictions, callees_sort_key=(), callees_filter=())
 
       This method for the :class:`Stats` class prints a list of all function
       that were called by the indicated function.  Aside from this reversal of
-      direction of calls (re: called vs was called by), the arguments and
-      ordering are identical to the :meth:`~pstats.Stats.print_callers` method.
+      direction of calls (re: called vs was called by), and names of the
+      keyword arguments, the arguments and ordering are identical to the
+      :meth:`~pstats.Stats.print_callers` method.
 
 
    .. method:: get_stats_profile()

Lib/pstats.py

@@ -232,10 +232,7 @@ def get_sort_arg_defs(self):
                 del dict[word]
         return self.sort_arg_dict
 
-    def sort_stats(self, *field):
-        if not field:
-            self.fcn_list = 0
-            return self
+    def get_sort_tuple_and_type(self, *field):
         if len(field) == 1 and isinstance(field[0], int):
             # Be compatible with old profiler
             field = [ {-1: "stdname",
@@ -249,26 +246,35 @@ def sort_stats(self, *field):
 
         sort_arg_defs = self.get_sort_arg_defs()
 
-        sort_tuple = ()
-        self.sort_type = ""
-        connector = ""
+        sort_fields = []
+        sort_field_types = []
         for word in field:
             if isinstance(word, SortKey):
                 word = word.value
-            sort_tuple = sort_tuple + sort_arg_defs[word][0]
-            self.sort_type += connector + sort_arg_defs[word][1]
-            connector = ", "
+            sort_fields.extend(sort_arg_defs[word][0])
+            sort_field_types.append(sort_arg_defs[word][1])
+
+        return (tuple(sort_fields), ", ".join(sort_field_types))
+
+    def sort_stats(self, *field):
+        if not field:
+            self.fcn_list = 0
+            return self
+
+        sort_tuple, self.sort_type = self.get_sort_tuple_and_type(*field)
 
         stats_list = []
         for func, (cc, nc, tt, ct, callers) in self.stats.items():
+            # matches sort_arg_defs: cc, nc, tt, ct, module, line, name, stdname
             stats_list.append((cc, nc, tt, ct) + func +
-                              (func_std_string(func), func))
+                              (func_std_string(func),))
 
         stats_list.sort(key=cmp_to_key(TupleComp(sort_tuple).compare))
 
+        # store only the order of the funcs
         self.fcn_list = fcn_list = []
         for tuple in stats_list:
-            fcn_list.append(tuple[-1])
+            fcn_list.append(tuple[4:7])
         return self
 
     def reverse_order(self):
@@ -432,28 +438,50 @@ def print_stats(self, *amount):
             print(file=self.stream)
         return self
 
-    def print_callees(self, *amount):
+    def print_callees(self, *amount, callees_sort_key=None, callees_filter=()):
         width, list = self.get_print_list(amount)
         if list:
             self.calc_callees()
 
+            sort_tuple = None
+            sort_type = None
+            if callees_sort_key:
+                if isinstance(callees_sort_key, str):
+                    callees_sort_key = (callees_sort_key,)
+                sort_tuple, sort_type = self.get_sort_tuple_and_type(*callees_sort_key)
+                print(f"   Callees ordered by: {sort_type}", file=self.stream)
+                print(file=self.stream)
+            if not isinstance(callees_filter, tuple):
+                callees_filter = (callees_filter,)
             self.print_call_heading(width, "called...")
             for func in list:
                 if func in self.all_callees:
-                    self.print_call_line(width, func, self.all_callees[func])
+                    self.print_call_line(width, func, self.all_callees[func],
+                                         sort_tuple=sort_tuple, sort_type=sort_type, sel_list=callees_filter)
                 else:
                     self.print_call_line(width, func, {})
             print(file=self.stream)
             print(file=self.stream)
         return self
 
-    def print_callers(self, *amount):
+    def print_callers(self, *amount, callers_sort_key=None, callers_filter=()):
         width, list = self.get_print_list(amount)
         if list:
+            sort_tuple = None
+            sort_type = None
+            if callers_sort_key:
+                if isinstance(callers_sort_key, str):
+                    callers_sort_key = (callers_sort_key,)
+                sort_tuple, sort_type = self.get_sort_tuple_and_type(*callers_sort_key)
+                print(f"   Callers ordered by: {sort_type}", file=self.stream)
+                print(file=self.stream)
+            if not isinstance(callers_filter, tuple):
+                callers_filter = (callers_filter,)
             self.print_call_heading(width, "was called by...")
             for func in list:
                 cc, nc, tt, ct, callers = self.stats[func]
-                self.print_call_line(width, func, callers, "<-")
+                self.print_call_line(width, func, callers, arrow="<-",
+                                     sort_tuple=sort_tuple, sort_type=sort_type, sel_list=callers_filter)
             print(file=self.stream)
             print(file=self.stream)
         return self
@@ -470,14 +498,41 @@ def print_call_heading(self, name_size, column_title):
         if subheader:
             print(" "*name_size + "    ncalls  tottime  cumtime", file=self.stream)
 
-    def print_call_line(self, name_size, source, call_dict, arrow="->"):
+    def print_call_line(self, name_size, source, call_dict, arrow="->", sort_tuple=None, sort_type=None, sel_list=()):
         print(func_std_string(source).ljust(name_size) + arrow, end=' ', file=self.stream)
         if not call_dict:
             print(file=self.stream)
             return
-        clist = sorted(call_dict.keys())
+
+        if sort_tuple:
+            stats_list = []
+            calls_only = False
+            for func, value in call_dict.items():
+                if isinstance(value, tuple):
+                    nc, cc, tt, ct = value  # cProfile orders it this way
+                    # matches sort_arg_defs: cc, nc, tt, ct, module, line, name, stdname
+                    stats_list.append((cc, nc, tt, ct) + func +
+                                      (func_std_string(func),))
+                else:
+                    if not calls_only and "time" in sort_type:
+                        raise TypeError("Caller/callee stats for %s do not have time information. "
+                                        "Try using cProfile instead of profile if you wish to record time by caller/callee."
+                                        % func_std_string(func))
+                    calls_only = True
+                    stats_list.append((None, value, None, None) + func +
+                                      (func_std_string(func),))
+
+            stats_list.sort(key=cmp_to_key(TupleComp(sort_tuple).compare))
+            funclist = [t[4:7] for t in stats_list]
+        else:
+            funclist = list(sorted(call_dict.keys()))
+
+        msg = ""
+        for selection in sel_list:
+            funclist, msg = self.eval_print_amount(selection, funclist, msg)
+
         indent = ""
-        for func in clist:
+        for func in funclist:
             name = func_std_string(func)
             value = call_dict[func]
             if isinstance(value, tuple):
@@ -494,6 +549,8 @@ def print_call_line(self, name_size, source, call_dict, arrow="->"):
                 left_width = name_size + 3
             print(indent*left_width + substats, file=self.stream)
             indent = " "
+        if msg:
+            print(msg, file=self.stream)
 
     def print_title(self):
         print('   ncalls  tottime  percall  cumtime  percall', end=' ', file=self.stream)

Lib/test/test_pstats.py

@@ -7,9 +7,18 @@
 
 import os
 import pstats
+import re
 import tempfile
 import cProfile
 
+_CALLERS_LINE = re.compile(r'''
+        (?P<ncalls>\d+) \s+
+        (?P<tottime>\d+(?:[.]\d+)?) \s+
+        (?P<cumtime>\d+(?:[.]\d+)?) \s+
+        (?P<name>(?P<filename>\w+[.]\w+):(?P<line>\d+)\((?P<func>\w+)\)
+         | \{built-in method (?P<builtin>\w+)\})$
+    ''', re.M|re.X)
+
 class AddCallersTestCase(unittest.TestCase):
     """Tests for pstats.add_callers helper."""
 
@@ -30,8 +39,8 @@ def test_combine_results(self):
 
 class StatsTestCase(unittest.TestCase):
     def setUp(self):
-        stats_file = support.findfile('pstats.pck')
-        self.stats = pstats.Stats(stats_file)
+        self.stats_file = support.findfile('pstats.pck')
+        self.stats = pstats.Stats(self.stats_file)
 
     def test_add(self):
         stream = StringIO()
@@ -65,6 +74,35 @@ def test_loading_wrong_types(self):
         with self.assertRaises(TypeError):
             stats.load_stats(42)
 
+    def test_print_callers_sorted(self):
+        stream = StringIO()
+        stats = pstats.Stats(self.stats_file, stream=stream)
+        stats.strip_dirs()
+        stats.sort_stats('ncalls')
+        # Sorted filtered callers
+        stats.print_callers(1, callers_sort_key='ncalls', callers_filter=5)
+        res = stream.getvalue()
+        lines = list(_CALLERS_LINE.finditer(res))
+        calls = [int(line.group('ncalls')) for line in lines]
+        self.assertLessEqual(len(lines), 5)
+        self.assertEqual(calls, sorted(calls, reverse=True))
+
+    def test_print_callers_unsorted(self):
+        stream = StringIO()
+        stats = pstats.Stats(self.stats_file, stream=stream)
+        stats.strip_dirs()
+        stats.sort_stats('ncalls')
+        # Unsorted unfiltered callers
+        stats.print_callers(1)
+        res = stream.getvalue()
+        lines = list(_CALLERS_LINE.finditer(res))
+        max_calls = max(int(line.group('ncalls')) for line in lines)
+        filenames = [line.group('filename') for line in lines]
+        # line numbers are sorted as ints, so just compare filenames
+        self.assertSequenceEqual(filenames, sorted(filenames))
+        self.assertGreater(len(lines), 5)
+        self.assertLess(int(lines[0].group('ncalls')), max_calls)
+
     def test_sort_stats_int(self):
         valid_args = {-1: 'stdname',
                       0: 'calls',

Misc/NEWS.d/next/Library/2024-10-22-19-36-47.gh-issue-125855.942kuz.rst

@@ -0,0 +1 @@
+Add new keyword arguments for sorting and filtering callers and callees to :func:`~pstats.Stats.print_callers` and :func:`~pstats.Stats.print_callees`. Contributed by Benjamin S Wolf.