File tree Expand file tree Collapse file tree 2 files changed +44
-1
lines changed Expand file tree Collapse file tree 2 files changed +44
-1
lines changed Original file line number Diff line number Diff line change 66from _colorize import ANSIColors
77
88from .pstats_collector import PstatsCollector
9+ from .stack_collectors import CollapsedStackCollector
910
1011
1112class SampleProfiler :
@@ -255,6 +256,9 @@ def sample(
255256 match output_format :
256257 case "pstats" :
257258 collector = PstatsCollector (sample_interval_usec )
259+ case "collapsed" :
260+ collector = CollapsedStackCollector ()
261+ filename = filename or f"collapsed.{ pid } .txt"
258262 case _:
259263 raise ValueError (f"Invalid output format: { output_format } " )
260264
@@ -281,6 +285,8 @@ def main():
281285 "The default sort is by cumulative time (--sort-cumulative)."
282286 "Format descriptions:\n "
283287 " pstats Standard Python profiler output format\n "
288+ " collapsed Stack traces in collapsed format (file:function:line;file:function:line;... count)\n "
289+ " Useful for generating flamegraphs with tools like flamegraph.pl"
284290 ),
285291 formatter_class = argparse .RawDescriptionHelpFormatter ,
286292 color = True ,
@@ -325,7 +331,7 @@ def main():
325331 )
326332 parser .add_argument (
327333 "--format" ,
328- choices = ["pstats" ],
334+ choices = ["pstats" , "collapsed" ],
329335 default = "pstats" ,
330336 help = "Output format (default: pstats)" ,
331337 )
Original file line number Diff line number Diff line change 1+ import collections
2+ import os
3+
4+ from .collector import Collector
5+
6+
7+ class StackTraceCollector (Collector ):
8+ def __init__ (self ):
9+ self .call_trees = []
10+ self .function_samples = collections .defaultdict (int )
11+
12+ def collect (self , stack_frames ):
13+ for thread_id , frames in stack_frames :
14+ if frames and len (frames ) > 0 :
15+ # Store the complete call stack (reverse order - root first)
16+ call_tree = list (reversed (frames ))
17+ self .call_trees .append (call_tree )
18+
19+ # Count samples per function
20+ for frame in frames :
21+ self .function_samples [frame ] += 1
22+
23+
24+ class CollapsedStackCollector (StackTraceCollector ):
25+ def export (self , filename ):
26+ stack_counter = collections .Counter ()
27+ for call_tree in self .call_trees :
28+ # Call tree is already in root->leaf order
29+ stack_str = ";" .join (
30+ f"{ os .path .basename (f [0 ])} :{ f [2 ]} :{ f [1 ]} " for f in call_tree
31+ )
32+ stack_counter [stack_str ] += 1
33+
34+ with open (filename , "w" ) as f :
35+ for stack , count in stack_counter .items ():
36+ f .write (f"{ stack } { count } \n " )
37+ print (f"Collapsed stack output written to { filename } " )
You can’t perform that action at this time.
0 commit comments