sharktop

#!/usr/bin/python3
# Based on example by https://gist.github.com/claymcleod


import os
import re
import time
import curses
import argparse
import threading
import subprocess


class SharkTop(object):
    # CONSTANTS
    PANEL_FPS = 0
    PANEL_RAW = 1
    PANEL_QUEUE = 2
    BAR_NORMAL = 0
    BAR_FILTER = 1

    # Tracer patterns
    PATTERN_FPS = re.compile(
        r"(.*)(framerate)(.*)(pad=\(string\))"
        r"(.*)(,)(.*)(fps=\(uint\))(\d+)"
    )
    PATTERN_QUEUE_LEVEL = re.compile(
        r"(.*)(queuelevel,)(.*queue=\(string\))(.*?)(,.*)"
        r"(size_buffers=\(uint\))(.*?)(,.*)(max_size_buffers=\(uint\))"
        r"(.*?)(,.*)"
    )
    PATTERN_ANSI = re.compile(
        r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])'
    )

    def __init__(self):
        # Properties
        self.args = None
        self.start_time = None

        # Filters
        self.pattern_string_fps = None
        self.pattern_string_queue = None
        self.filter_pattern_fps = None
        self.filter_pattern_queue = None

        # State
        self.close = False
        self.bar = self.BAR_NORMAL
        self.panel = self.PANEL_FPS

        # Data
        self.raw = []
        self.framerate = {}
        self.queue_level = {}

    def _update_values(self, process, args):
        self.filter_pattern_fps = re.compile(self.pattern_string_fps)
        self.filter_pattern_queue = \
            re.compile(self.pattern_string_queue)

        for line in iter(process.stdout.readline, ''):
            if line == '':
                self._add_to_raw_output(
                    "\n#### End of raw stream ####\n"
                )
                break

            line = line.decode("utf-8")

            line = self.PATTERN_ANSI.sub('', line)

            self._add_to_raw_output(line)

            not self._parse_fps(line) and not self._parse_queue(line)

    def _add_to_raw_output(self, text):
        self.raw.append(text)
        self.raw = self.raw[-200:]

    def _parse_strace(self, line):
        pruned_line = None
        search_strace = self.PATTERN_STRACE.search(line)
        if search_strace is not None:
            pruned_line = search_strace.group(1)
        
        return pruned_line

    def _parse_fps(self, line):
        matched = False

        search_fps = self.PATTERN_FPS.search(
            line
        )
        if search_fps is not None and \
                self.filter_pattern_fps.search(
                    search_fps.group(5)) is not None:
            self.framerate[search_fps.group(5)] = search_fps.group(9)
            matched = True

        return matched

    def _parse_queue(self, line):
        matched = False

        search_queue = self.PATTERN_QUEUE_LEVEL.search(
            line
        )
        if search_queue is not None and \
                self.filter_pattern_queue.search(
                    search_queue.group(4)) is not None and \
                search_queue.group(4) != "%s":
            self.queue_level[search_queue.group(4)] = (
                search_queue.group(7), search_queue.group(10)
            )
            matched = True

        return matched

    def _get_arguments(self):
        parser = argparse.ArgumentParser()
        group = parser.add_mutually_exclusive_group(required=True)
        group.add_argument(
            '-p', '--pid',
            help='PID of the GstShark enabled process',
            type=int,
            default=-1
        )
        group.add_argument(
            '-i', '--pipeline',
            type=str, nargs='+',
            help='Command to run GStreamer pipeline',
            default=None
        )
        parser.add_argument(
            '-r', '--referesh_rate',
            help='Update rate per second',
            type=int,
            default=2
        )
        parser.add_argument(
            '-f', '--filter',
            help='Pattern to match for the output',
            type=str,
            default=r".*"
        )
        args = parser.parse_args()

        self.pattern_string_fps = args.filter
        self.pattern_string_queue = args.filter

        return args

    def _manage_panels(self, pressed_key, cursor_y, cursor_y_content):
        if pressed_key == ord('0'):
            self.panel = self.PANEL_FPS
            cursor_y = 2
            cursor_y_content = 0
        elif pressed_key == ord('1'):
            self.panel = self.PANEL_QUEUE
            cursor_y = 2
            cursor_y_content = 0
        elif pressed_key == ord('2'):
            self.panel = self.PANEL_RAW
            cursor_y = 2
            cursor_y_content = 0
        elif pressed_key == ord('f') and self.panel != self.PANEL_RAW:
            self.bar = self.BAR_FILTER
            if self.panel == self.PANEL_FPS:
                self.pattern_string_fps = ""
            elif self.panel == self.PANEL_QUEUE:
                self.pattern_string_queue = ""

        return cursor_y, cursor_y_content

    def _manage_dynamic_filter(self, pressed_key):
        if pressed_key != -1:
            if self.panel == self.PANEL_FPS:
                if pressed_key == curses.KEY_ENTER or \
                        pressed_key == 10 or pressed_key == 13 or \
                        pressed_key == 27:
                    self.bar = self.BAR_NORMAL
                    if self.pattern_string_fps.strip() == '':
                        self.pattern_string_fps = ".*"
                else:
                    if pressed_key == curses.KEY_BACKSPACE:
                        if len(self.pattern_string_fps) > 0:
                            self.pattern_string_fps = \
                                self.pattern_string_fps[:-1]
                    else:
                        self.pattern_string_fps += chr(pressed_key)

                    self.filter_pattern_fps = re.compile(
                        self.pattern_string_fps
                    )
                    self.framerate.clear()
            elif self.panel == self.PANEL_QUEUE:
                if pressed_key == curses.KEY_ENTER or \
                        pressed_key == 10 or pressed_key == 13 or \
                        pressed_key == 27:
                    self.bar = self.BAR_NORMAL
                    if self.pattern_string_queue.strip() == '':
                        self.pattern_string_queue = ".*"
                else:
                    if pressed_key == curses.KEY_BACKSPACE:
                        if len(self.pattern_string_queue) > 0:
                            self.pattern_string_queue = \
                                self.pattern_string_queue[:-1]
                    else:
                        self.pattern_string_queue += chr(pressed_key)

                    self.filter_pattern_queue = \
                        re.compile(self.pattern_string_queue)
                    self.queue_level.clear()

    def _update_cursor_fps(self, pressed_key, cursor_y,
            cursor_y_content, height):
        if pressed_key == curses.KEY_DOWN:
            cursor_y_content = min(
                len(self.framerate.keys()),
                max(0, cursor_y_content + 1)
            )
            cursor_y = min(height - 2, max(2, cursor_y + 1))
        elif pressed_key == curses.KEY_UP:
            cursor_y_content = min(
                len(self.framerate.keys()),
                max(0, cursor_y_content - 1)
            )
            if cursor_y_content <= height - 5:
                cursor_y = min(height - 2, max(2, cursor_y - 1))

        return cursor_y, cursor_y_content

    def _render_panel_fps(self, stdscr, pressed_key,
            cursor_y, cursor_y_content, width, height):
        self._render_title_fps(stdscr, width)

        self._render_headers_fps(stdscr, width)

        self._render_values_fps(
            stdscr, cursor_y,
            cursor_y_content, width, height
        )

        self._render_statusbar_fps(stdscr, width, height)

    def _render_title_fps(self, stdscr, width):
        stdscr.attron(curses.color_pair(4))
        if self.args.pid == -1:
            stdscr.addstr(0, 0, "Pipeline:")
            stdscr.addstr(
                0, 9, " ".join(self.args.pipeline)[:width - 10]
            )
        else:
            stdscr.addstr(0, 0, "PID:")
            stdscr.addstr(
                0, 4, " " + str(self.args.pid)[:width - 5]
            )
        stdscr.attroff(curses.color_pair(4))

    def _render_headers_fps(self, stdscr, width):
        stdscr.attron(curses.color_pair(1))
        stdscr.addstr(1, 0, "Element" + (" " * (width // 2 - 7)))
        stdscr.addstr(
            1, width // 2,
            "Framerate" + (" " * (width // 2 - 9))
        )
        stdscr.attroff(curses.color_pair(1))

    def _render_values_fps(self, stdscr, cursor_y,
            cursor_y_content, width, height):
        stdscr.attron(curses.color_pair(2))

        offset = max(cursor_y_content - height + 3 + 1, 0)
        keys = list(self.framerate.keys())
        for i in range(2, min(height, len(keys) - offset + 2)):
            key = keys[i - 2 + offset]

            if i == cursor_y:
                stdscr.attroff(curses.color_pair(2))
                stdscr.attron(curses.color_pair(1))

            stdscr.addstr(i, 0, key + (" " * (width // 2 - len(key))))
            stdscr.addstr(
                i, width // 2,
                str(self.framerate[key]) +
                (" " * (width // 2 - len(str(self.framerate[key])) - 1))
            )

            if i == cursor_y:
                stdscr.attroff(curses.color_pair(1))
                stdscr.attron(curses.color_pair(2))

        stdscr.attroff(curses.color_pair(2))

    def _render_statusbar_fps(self, stdscr, width, height):
        statusbarstr = (
            "'q' exit | '0' FPS | '1' Queue | '2' Raw" +
            " | uptime: {:.2f}s | 'f' filter: {}"
        ).format(
            time.time() - self.start_time,
            self.pattern_string_fps
        )[-1 * width + 1:]
        stdscr.attron(curses.color_pair(3))
        stdscr.addstr(height-1, 0, statusbarstr)
        stdscr.addstr(
            height-1, len(statusbarstr),
            " " * (width - len(statusbarstr) - 1)
        )
        stdscr.attroff(curses.color_pair(3))

    def _update_cursor_queue(self, pressed_key, cursor_y,
            cursor_y_content, height):
        if pressed_key == curses.KEY_DOWN:
            cursor_y_content = min(
                len(self.queue_level.keys()),
                max(0, cursor_y_content + 1)
            )
            cursor_y = min(height - 2, max(2, cursor_y + 1))
        elif pressed_key == curses.KEY_UP:
            cursor_y_content = min(
                len(self.queue_level.keys()),
                max(0, cursor_y_content - 1)
            )
            if cursor_y_content <= height - 5:
                cursor_y = min(height - 2, max(2, cursor_y - 1))

        return cursor_y, cursor_y_content

    def _render_panel_queue(self, stdscr, pressed_key,
            cursor_y, cursor_y_content, width, height):
        self._render_title_queue(stdscr, width)

        self._render_headers_queue(stdscr, width)

        self._render_values_queue(
            stdscr, cursor_y,
            cursor_y_content, width, height
        )

        self._render_statusbar_queue(stdscr, width, height)

    def _render_title_queue(self, stdscr, width):
        stdscr.attron(curses.color_pair(4))
        if self.args.pid == -1:
            stdscr.addstr(0, 0, "Pipeline:")
            stdscr.addstr(
                0, 9, " ".join(self.args.pipeline)[:width - 10]
            )
        else:
            stdscr.addstr(0, 0, "PID:")
            stdscr.addstr(
                0, 4, " " + str(self.args.pid)[:width - 5]
            )
        stdscr.attroff(curses.color_pair(4))

    def _render_headers_queue(self, stdscr, width):
        stdscr.attron(curses.color_pair(1))
        stdscr.addstr(1, 0, "Element" + (" " * (width // 2 - 7)))
        stdscr.addstr(
            1, width // 2,
            "QueueLevel" + (" " * (width // 2 - 10))
        )
        stdscr.attroff(curses.color_pair(1))

    def _render_values_queue(self, stdscr, cursor_y,
            cursor_y_content, width, height):
        stdscr.attron(curses.color_pair(2))

        offset = max(cursor_y_content - height + 3 + 1, 0)
        keys = list(self.queue_level.keys())
        for i in range(2, min(height, len(keys) - offset + 2)):
            key = keys[i - 2 + offset]

            if i == cursor_y:
                stdscr.attroff(curses.color_pair(2))
                stdscr.attron(curses.color_pair(1))

            stdscr.addstr(i, 0, key + (" " * (width // 2 - len(key))))
            text = \
                self.queue_level[key][0] + "/" + \
                self.queue_level[key][1]
            stdscr.addstr(
                i, width // 2,
                text +
                (" " * (width // 2 - len(text) - 1))
            )

            if i == cursor_y:
                stdscr.attroff(curses.color_pair(1))
                stdscr.attron(curses.color_pair(2))

        stdscr.attroff(curses.color_pair(2))

    def _render_statusbar_queue(self, stdscr, width, height):
        statusbarstr = (
            "'q' exit | '0' FPS | '1' Queue | '2' Raw" +
            " | uptime: {:.2f}s | 'f' filter: {}"
        ).format(
            time.time() - self.start_time,
            self.pattern_string_queue
        )[-1 * width + 1:]
        stdscr.attron(curses.color_pair(3))
        stdscr.addstr(height-1, 0, statusbarstr)
        stdscr.addstr(
            height-1, len(statusbarstr),
            " " * (width - len(statusbarstr) - 1)
        )
        stdscr.attroff(curses.color_pair(3))

    def _render_panel_raw(self, stdscr, width, height):
        # Render title
        stdscr.attron(curses.color_pair(4))
        if self.args.pid == -1:
            stdscr.addstr(0, 0, "Pipeline:")
            stdscr.addstr(
                0, 9, " ".join(self.args.pipeline)[:width - 10]
            )
        else:
            stdscr.addstr(0, 0, "PID:")
            stdscr.addstr(
                0, 4, " " + str(self.args.pid)[:width - 5]
            )
        stdscr.attroff(curses.color_pair(4))

        # Render raw outputs
        for i in range(min(height - 2, len(self.raw))):
            text = self.raw[-1 * min((height - 2 - i), len(self.raw))]
            stdscr.addstr(i + 1, 0, text[:width])

        # Render status bar
        statusbarstr = (
            "'q' exit | '0' FPS | '1' Queue | '2' Raw" +
            " | uptime: {:.2f}s"
        ).format(
            time.time() - self.start_time
        )
        stdscr.attron(curses.color_pair(3))
        stdscr.addstr(height-1, 0, statusbarstr)
        stdscr.addstr(
            height-1, len(statusbarstr),
            " " * (width - len(statusbarstr) - 1)
        )
        stdscr.attroff(curses.color_pair(3))

    def _clear_refresh(self, stdscr):
        stdscr.clear()
        stdscr.refresh()

    def _initialize_colors(self):
        curses.start_color()
        curses.init_pair(1, curses.COLOR_CYAN, curses.COLOR_WHITE)
        curses.init_pair(2, curses.COLOR_BLUE, curses.COLOR_BLACK)
        curses.init_pair(3, curses.COLOR_BLACK, curses.COLOR_WHITE)
        curses.init_pair(4, curses.COLOR_RED, curses.COLOR_YELLOW)

    def _initialize_window(self, stdscr):
        stdscr.clear()
        height, width = stdscr.getmaxyx()

        return height, width

    def _update_window(self, stdscr):
        stdscr.refresh()

        # Wait for next input
        stdscr.timeout(
            int(min(1 / self.args.referesh_rate, .5) * 10e2)
        )
        pressed_key = stdscr.getch()
        curses.napms(int(1000 / 120))

        return pressed_key

    def _stop_update_thread(self):
        self.close = True

    def draw_menu(self, stdscr):
        curses.curs_set(0)

        pressed_key = 0
        cursor_y = 2
        cursor_y_content = 0

        self._clear_refresh(stdscr)
        self._initialize_colors()

        while not (pressed_key == ord('q') and
                   self.bar == self.BAR_NORMAL):
            height, width = self._initialize_window(stdscr)

            if self.bar == self.BAR_NORMAL:
                cursor_y, cursor_y_content = self._manage_panels(
                    pressed_key, cursor_y, cursor_y_content
                )
            else:
                self._manage_dynamic_filter(pressed_key)

            if self.panel == self.PANEL_FPS:
                cursor_y, cursor_y_content = self._update_cursor_fps(
                    pressed_key, cursor_y, cursor_y_content, height
                )
                self._render_panel_fps(
                    stdscr, pressed_key, cursor_y, cursor_y_content,
                    width, height
                )
            if self.panel == self.PANEL_QUEUE:
                cursor_y, cursor_y_content = self._update_cursor_queue(
                    pressed_key, cursor_y, cursor_y_content, height
                )
                self._render_panel_queue(
                    stdscr, pressed_key, cursor_y, cursor_y_content,
                    width, height
                )
            elif self.panel == self.PANEL_RAW:
                self._render_panel_raw(stdscr, width, height)

            pressed_key = self._update_window(stdscr)

        self._stop_update_thread()

    def _initialize_update_thread(self, process, args):
        thread = threading.Thread(
            target=self._update_values,
            args=(process, args,)
        )
        thread.daemon = True
        thread.start()

    def _run_command(self, args):
        env = os.environ.copy()
        env["GST_DEBUG"] = "GST_TRACER:7"
        env["GST_TRACERS"] = "framerate;queuelevel"
        process = subprocess.Popen(
            args.pipeline, stdout=subprocess.PIPE, env=env,
            stderr=subprocess.STDOUT
        )

        self.start_time = time.time()

        return process

    def _attach_pid(self, args):
        command = f"peekfd -c -t -d -n {args.pid} 0 1 2"
        process = subprocess.Popen(
            command.split(), stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT
        )

        self.start_time = time.time()

        return process

    def start(self):
        self.args = self._get_arguments()

        process = None
        if self.args.pid != -1:
            process = self._attach_pid(self.args)
        else:
            process = self._run_command(self.args)

        self._initialize_update_thread(process, self.args)

        curses.wrapper(self.draw_menu)

        process.kill()


if __name__ == "__main__":
    sharktop = SharkTop()
    sharktop.start()