Process tree using psutil

pst-psutil.py

import psutil
import sys


COLOR_RESET = "\033[0m"
COLOR_LINE = "\033[1;30m"
COLOR_NAME = "\033[1;36m"
COLOR_PID = "\033[1;35m"
COLOR_USER = "\033[1;34m"
COLOR_CMD = "\033[0;37m"


def find_roots(procs):
    pids = set(proc['pid'] for proc in procs)
    return [proc['pid'] for proc in procs if proc['ppid'] not in pids or proc['pid'] == 1]


def read_procs(attrs=['pid', 'ppid', 'name', 'username', 'cmdline']):
    return [proc.info for proc in psutil.process_iter(attrs)]


def show_tree(start, procs):
    def show_children(depth, more_at_depth, parent, last_username):
        children = (proc for proc in procs if proc['ppid'] == parent)
        depth += 1
        items = sorted(children, key=lambda proc: (proc['name'].lower(), proc['pid']))
        if items:
            last = items[-1]
            rest = items[:-1]
            for item in rest:
                more = more_at_depth.copy()
                more[depth] = True
                show_proc(depth, False, more, last_username, item)
            more = more_at_depth.copy()
            more[depth] = False
            show_proc(depth, True, more, last_username, last)

    def show_proc(depth, last, more_at_depth, last_username, proc):
        if depth != 0:
            sys.stdout.write(COLOR_LINE)
            for i in range(1, depth + 1):
                more = more_at_depth[i]
                sys.stdout.write('  ')
                if i == depth:
                    sys.stdout.write('└─' if last else '├─')
                else:
                    sys.stdout.write('│ ' if more else '  ')
            sys.stdout.write(COLOR_RESET)

        name = f"{COLOR_NAME}{proc['name']}{COLOR_RESET}"
        pid = f"{COLOR_PID}{proc['pid']}{COLOR_RESET}"
        user = (f"{COLOR_USER}{proc['username']}{COLOR_RESET}"
                if proc['username'] and proc['username'] != last_username else '')
        cmd = (f"{COLOR_CMD}{' '.join(proc['cmdline'][1:])}{COLOR_RESET}"
               if proc['cmdline'] and len(proc['cmdline']) > 1 else '')

        sys.stdout.write(f"{name},{pid}{',' + user if user else ''}{' ' + cmd if cmd else ''}\n")
        show_children(depth, more_at_depth, proc['pid'], proc['username'])

    root = [proc for proc in procs if proc['pid'] == start][0]
    show_proc(0, True, {}, 'root', root)


def main():
    sys.stdout.reconfigure(encoding='utf-8')
    procs = read_procs()
    for root in find_roots(procs):
        show_tree(root, procs)


if __name__ == "__main__":
    main()