exclude all processes after self (fix bug #24) (#28)

* exclude all processes after self

* remove own and descendants after having all processes in a tree

* rename to processes

* processes dict in get_all

* extract remove own process

* remove process takes pid as argument

* conform flake8

* remove process from parent

* assert all parents and children in processes (no orphans)

px/px_process.py

@@ -222,46 +222,44 @@ def resolve_links(processes, now):
     Also, all processes will have a (possibly empty) "children" field containing
     a set of references to child processes.
     """
-    pid2process = {}
-    for process in processes:
-        # Guard against duplicate PIDs
-        assert process.pid not in pid2process
-
-        pid2process[process.pid] = process
-
+    for process in processes.values():
         process.children = set()
 
-    if 0 not in pid2process:
+    if 0 not in processes:
         kernel_process = create_kernel_process(now)
+        processes[0] = kernel_process
 
-        processes.append(kernel_process)
-        pid2process[0] = kernel_process
-
-    for process in processes:
+    for process in processes.values():
         if process.pid == 0:
             process.parent = None
         else:
-            process.parent = pid2process[process.ppid]
+            process.parent = processes[process.ppid]
             process.parent.children.add(process)
 
 
+def remove_process_and_descendants(processes, pid):
+    process = processes[pid]
+    process.parent.children.remove(process)
+    toexclude = [process]
+    while toexclude:
+        process = toexclude.pop()
+        del processes[process.pid]
+        for child in process.children:
+            toexclude.append(child)
+
+
 def get_all():
-    all = []
+    processes = {}
     ps_lines = call_ps()
     now = datetime.datetime.now().replace(tzinfo=dateutil.tz.tzlocal())
     for ps_line in ps_lines:
         process = ps_line_to_process(ps_line, now)
-        if process.pid == os.getpid():
-            # Finding ourselves is just confusing
-            continue
-        if process.ppid == os.getpid():
-            # Finding the ps we spawned is also confusing
-            continue
-        all.append(process)
+        processes[process.pid] = process
 
-    resolve_links(all, now)
+    resolve_links(processes, now)
+    remove_process_and_descendants(processes, os.getpid())
 
-    return all
+    return processes.values()
 
 
 def order_best_last(processes):

tests/px_process_test.py

@@ -108,9 +108,11 @@ def _validate_references(processes):
 
         assert type(process.children) is set
         if process.parent:
+            assert process.parent in processes
             assert process in process.parent.children
 
         for child in process.children:
+            assert child in processes
             assert child.parent == process