Monday, April 7, 2014

[Level 2] pstree implementation with Python.

Because I couldn't find pstree command in our company's product. Therefore, I just search from internet and found someone implement pstree layout with Python. (http://stackoverflow.com/questions/16395207/create-a-process-tree-like-pstree-command-with-python-in-linux) And I just modify it and create a pstree.py system utility.
#!/bin/env python
'''
## Implement pstree in Python

## data structure of tree, cmd_list
tree = {
  0: [1],
  1: [2, 3],
  2: [5, 6, 7, 8],
  ...
}

cmd_list = {
  0: '/sbin/init',
  1: '[kthreadd]',
  ...
}
'''

import os
import sys
import re
tree = {}
cmd_list = {}
def printTree(parent, tree, cmd_list, indent=''):
  print '%s:%s' % (parent, cmd_list[parent])
  if parent not in tree:
    return
  for child in tree[parent][:-1]:
    sys.stdout.write(indent + '|-')
    printTree(child, tree, cmd_list, indent + '| ')
  child = tree[parent][-1]
  sys.stdout.write(indent + '`-')
  printTree(child, tree, cmd_list, indent + '  ')

## __main__
for l in [ re.sub('[ ]+', ' ', l.lstrip().rstrip()) for l in os.popen('ps -e -o pid,ppid,command --no-heading').readlines() ]:
    pid = int(l.split(' ')[0])
    ppid = int(l.split(' ')[1])
    cmd = ' '.join(l.split(' ')[2:])
    if ppid not in tree.keys():
        tree[ppid] = []
    tree[ppid].append(pid)
    cmd_list[pid] = cmd

# print tree
# print cmd_list
pid = 1
if len(sys.argv) > 1:
  pid = int(sys.argv[1])
printTree(pid, tree, cmd_list)

Wish this helps.

regards,
Stanley Huang