feat: adds a --watch flag to only showing stats that contain a specific statement

profiling/__main__.py

@@ -176,10 +176,10 @@ def get_title(src_name, src_type=None):
     return os.path.basename(src_name)
 
 
-def make_viewer(mono=False, *loop_args, **loop_kwargs):
+def make_viewer(mono=False, watch=None, *loop_args, **loop_kwargs):
     """Makes a :class:`profiling.viewer.StatisticsViewer` with common options.
     """
-    viewer = StatisticsViewer()
+    viewer = StatisticsViewer(watch=watch)
     loop = viewer.loop(*loop_args, **loop_kwargs)
     if mono:
         loop.screen.set_terminal_properties(1)
@@ -505,6 +505,7 @@ def wrapped(argv, module, command, **kwargs):
 
 viewer_options = Params([
     click.option('--mono', is_flag=True, help='Disable coloring.'),
+    click.option('--watch', help='Filter traces with a specific statement.'),
 ])
 onetime_profiler_options = Params([
     click.option(
@@ -536,7 +537,7 @@ def wrapped(argv, module, command, **kwargs):
 
 def __profile__(filename, code, globals_, profiler_factory,
                 pickle_protocol=remote.PICKLE_PROTOCOL, dump_filename=None,
-                mono=False):
+                mono=False, watch=None):
     frame = sys._getframe()
     profiler = profiler_factory(base_frame=frame, base_code=code)
     profiler.start()
@@ -552,7 +553,7 @@ def __profile__(filename, code, globals_, profiler_factory,
     profiler.stats.discard_child(frame.f_code)
     if dump_filename is None:
         try:
-            profiler.run_viewer(get_title(filename), mono=mono)
+            profiler.run_viewer(get_title(filename), mono=mono, watch=watch)
         except KeyboardInterrupt:
             pass
     else:
@@ -580,13 +581,13 @@ def collect_usage_pieces(self, ctx):
 @onetime_profiler_options
 @viewer_options
 def profile(script, argv, profiler_factory,
-            pickle_protocol, dump_filename, mono):
+            pickle_protocol, dump_filename, mono, watch):
     """Profile a Python script."""
     filename, code, globals_ = script
     sys.argv[:] = [filename] + list(argv)
     __profile__(filename, code, globals_, profiler_factory,
                 pickle_protocol=pickle_protocol, dump_filename=dump_filename,
-                mono=mono)
+                mono=mono, watch=watch)
 
 
 @cli.command('live-profile', aliases=['live'], cls=ProfilingCommand)
@@ -595,7 +596,7 @@ def profile(script, argv, profiler_factory,
 @live_profiler_options
 @viewer_options
 def live_profile(script, argv, profiler_factory, interval, spawn, signum,
-                 pickle_protocol, mono):
+                 pickle_protocol, mono, watch):
     """Profile a Python script continuously."""
     filename, code, globals_ = script
     sys.argv[:] = [filename] + list(argv)
@@ -605,7 +606,7 @@ def live_profile(script, argv, profiler_factory, interval, spawn, signum,
     if pid:
         # parent
         os.close(stderr_w_fd)
-        viewer, loop = make_viewer(mono)
+        viewer, loop = make_viewer(mono, watch)
         # loop.screen._term_output_file = open(os.devnull, 'w')
         title = get_title(filename)
         client = ProfilingClient(viewer, loop.event_loop, parent_sock, title)
@@ -704,11 +705,11 @@ def remote_profile(script, argv, profiler_factory, interval, spawn, signum,
 @click.argument('src', type=ViewerSource(),
                 default=config_default('endpoint', DEFAULT_ENDPOINT))
 @viewer_options
-def view(src, mono):
+def view(src, mono, watch):
     """Inspect statistics by TUI view."""
     src_type, src_name = src
     title = get_title(src_name, src_type)
-    viewer, loop = make_viewer(mono)
+    viewer, loop = make_viewer(mono, watch)
     if src_type == 'dump':
         time = datetime.fromtimestamp(os.path.getmtime(src_name))
         with open(src_name, 'rb') as f:
@@ -743,7 +744,7 @@ def view(src, mono):
 @viewer_options
 def timeit_profile(stmt, number, repeat, setup,
                    profiler_factory, pickle_protocol, dump_filename, mono,
-                   **_ignored):
+                   watch, **_ignored):
     """Profile a Python statement like timeit."""
     del _ignored
     globals_ = {}

profiling/profiler.py

@@ -85,17 +85,17 @@ def dump(self, dump_filename, pickle_protocol=pickle.HIGHEST_PROTOCOL):
         with open(dump_filename, 'wb') as f:
             pickle.dump((self.__class__, result), f, pickle_protocol)
 
-    def make_viewer(self, title=None, at=None):
+    def make_viewer(self, title=None, at=None, watch=None):
         """Makes a statistics viewer from the profiling result.
         """
-        viewer = StatisticsViewer()
+        viewer = StatisticsViewer(watch=watch)
         viewer.set_profiler_class(self.__class__)
         stats, cpu_time, wall_time = self.result()
         viewer.set_result(stats, cpu_time, wall_time, title=title, at=at)
         viewer.activate()
         return viewer
 
-    def run_viewer(self, title=None, at=None, mono=False,
+    def run_viewer(self, title=None, at=None, mono=False, watch=None,
                    *loop_args, **loop_kwargs):
         """A shorter form of:
 
@@ -106,7 +106,7 @@ def run_viewer(self, title=None, at=None, mono=False,
            loop.run()
 
         """
-        viewer = self.make_viewer(title, at=at)
+        viewer = self.make_viewer(title, at=at, watch=watch)
         loop = viewer.loop(*loop_args, **loop_kwargs)
         if mono:
             loop.screen.set_terminal_properties(1)

profiling/viewer.py

@@ -786,7 +786,8 @@ def unhandled_input(self, key):
         if key in ('q', 'Q'):
             raise urwid.ExitMainLoop()
 
-    def __init__(self):
+    def __init__(self, watch=None):
+        self.watch = watch
         self.table = StatisticsTable(self)
         self.widget = urwid.Padding(self.table, right=1)
 
@@ -803,8 +804,18 @@ def set_profiler_class(self, profiler_class):
         self.table = table_class(self)
         self.widget.original_widget = self.table
 
+    def is_watched(self, stat):
+        if stat.name == self.watch:
+            return True
+        for child in stat.children:
+            if self.is_watched(child):
+                return True
+
     def set_result(self, stats, cpu_time=0.0, wall_time=0.0,
                    title=None, at=None):
+        if self.watch and not self.is_watched(stats):
+            return
+
         self._final_result = (stats, cpu_time, wall_time, title, at)
         if not self.paused:
             self.update_result()