ps as tree view

pst

#!/usr/bin/env python
# usage:
#     $ pst
#     $ pst <pid>
#     $ pst <command>

import sys
import subprocess as sp
import re


def main():
    rootpid = []
    if len(sys.argv) >= 2:
        for pid in sys.argv[1:]:
            try:
                rootpid.append(int(pid))
            except ValueError:
                pass
            try:
                pids = sp.check_output(['pidof', pid]).strip()
                rootpid.extend([int(pid) for pid in pids.split()])
            except sp.CalledProcessError:
                pass

    if len(rootpid) == 0:
        sys.exit(sp.call(['ps', 'axjf']))
    p = sp.Popen(['ps', 'axjf'], stdout=sp.PIPE)
    reobj = re.compile(r'\s*(\d+)\s+(\d+)\s+')
    try:
        tasks = set(rootpid)
        print(p.stdout.readline().rstrip())  # caption
        for line in p.stdout:
            line = line.rstrip()
            mobj = reobj.search(line)
            if mobj is None:
                return
            ppid, pid = mobj.groups()
            try:
                ppid = int(ppid)
                pid = int(pid)
                if ppid in tasks or pid in tasks:
                    print(str(line))
                    tasks.add(pid)
            except ValueError:
                print('E: ' + str(line))
    finally:
        try:
            p.kill()
        except:
            pass


if __name__ == "__main__":
    main()

# vim: set ts=4 sw=4 et: -*- coding: utf-8 -*- filetype=python