diff --git a/profiling/__main__.py b/profiling/__main__.py index 418a9b6..dc03bb7 100644 --- a/profiling/__main__.py +++ b/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_ = {} diff --git a/profiling/profiler.py b/profiling/profiler.py index eee0e86..5e0e0b8 100644 --- a/profiling/profiler.py +++ b/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) diff --git a/profiling/viewer.py b/profiling/viewer.py index 36dbb1e..7c1ff14 100644 --- a/profiling/viewer.py +++ b/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()