Create a process tree like pstree command with python in linux

Question

I am new to python. I want to write a program outputting a tree-like drawing on the stdout. My ideal output is like:

0
|__0
|__4
|  |__360
|      |__1000
272
|__3460

The data I collected is as follows:

0       : [0, 4]
4       : [360]
272     : [3460]
368     : [4184]
472     : [504, 576, 7016]
568     : [584, 640]
576     : [664, 672]
640     : [1048]
664     : [368, 372, 512, 788]
788     : [2120, 2720, 2976, 2996, 3956, 3980]

The left column is parent process id, and right column is child process id. I put the data in a dictionary called dic. So the dictionary key is the parent process id, and the dictionary value is a list consisting of child process IDs.

My code is like this:

for key in dic.keys():
    print key, '\n|'
    for v in dic[key]:
        print '__', v, '\n|'

The question is I can only output two layer tree. Take the data for example, 576 as a parent id is also a child id of 472. So 576, 664, 672 should be placed in the sub-tree of 472. My code doesn't work for this. It seems that we need to use recursive functions. But I don't know how to handle it.

Could you guys give me hints?

---

EDIT: From the data I collected, there are some parent IDs which does not have grand parent. So the ultimate output should be a forest. Not a singled rooted tree.

Answer 1

How about this:

def printTree(parent, tree, indent=''):
  print parent
  if parent not in tree:
    return
  for child in tree[parent][:-1]:
    sys.stdout.write(indent + '|-')
    printTree(child, tree, indent + '| ')
  child = tree[parent][-1]
  sys.stdout.write(indent + '`-')
  printTree(child, tree, indent + '  ')

tree = {
  0       : [0, 4],
  4       : [360],
  272     : [3460],
  368     : [4184],
  472     : [504, 576, 7016],
  568     : [584, 640],
  576     : [664, 672],
  640     : [1048],
  664     : [368, 372, 512, 788],
  788     : [2120, 2720, 2976, 2996, 3956, 3980]
}

printTree(472, tree)

printTree(472, tree)
472
|-504
|-576
| |-664
| | |-368
| | | `-4184
| | |-372
| | |-512
| | `-788
| |   |-2120
| |   |-2720
| |   |-2976
| |   |-2996
| |   |-3956
| |   `-3980
| `-672
`-7016

Maybe that's how you like it, I don't know.

It does not have any checks built in for recursions, so if you try it on 0, it will run into an endless recursion (and abort eventually due to a stack overflow). You could check for recursions yourself by passing a trace of the already processed nodes.

This also does not find the list of tree roots in your forest, so you will have to do that as well. (But that sounds like another question.)