Merge pull request #17 from joknarf/threads

Add Threads display with -T option

.gitignore

@@ -4,3 +4,5 @@ dist
 pgtree.egg-info
 __pycache__
 *.pyc
+.coverage
+coverage.xml

.travis.yml

@@ -14,9 +14,11 @@ script:
   - python pgtree/pgtree.py
   - if [[ $TRAVIS_PYTHON_VERSION == 2.6 ]]; then exit 0; fi
   - if [[ $TRAVIS_PYTHON_VERSION == 2.7 ]]; then exit 0; fi
-  - pip install lint codecov incremental
+  - pip install lint coverage pytest pytest-cov incremental
   - pylint pgtree/pgtree.py ; echo done
   - python -m unittest discover -s tests/
-  - coverage run tests/test_pgtree.py
+  - pytest --cov
+  - curl -Os https://uploader.codecov.io/latest/linux/codecov
+  - chmod +x codecov
 after_success:
   - bash <(curl -s https://codecov.io/bash)

README.md

@@ -18,14 +18,16 @@ The code must be compatible with python 2.x + 3.x
 Should work on any Unix that can execute :
 ```
 # /usr/bin/pgrep 
-# /usr/bin/ps -e -o pid,ppid,stime,user,ucomm,args
+# /usr/bin/ps ax -o pid,ppid,stime,user,ucomm,args
 ```
 
 if `pgrep` command not available (AIX), pgtree uses built-in pgrep (`-f -i -x -u <user>` supported).
 
-_Tested on various versions of RedHat / CentOS / Ubuntu / Debian / Suse / MacOS / Solaris / AIX including old versions_
+`-T` option to display threads only works if `ps ax -T -o spid,ppid` available on system (ubuntu/redhat...)
 
-_(uses -o comm on Solaris)_
+_pgtree Tested on various versions of RedHat / CentOS / Ubuntu / Debian / Suse / FreeBSD / ArchLinux / MacOS / Solaris / AIX including old versions_
+
+_(uses -o fname on Solaris)_
 
 ## Installation
 FYI, the `pgtree/pgtree.py` is standalone and can be directly copied/used anywhere without any installation.
@@ -34,14 +36,7 @@ installation using pip:
 ```
 # pip install pgtree
 ```
-installation using setup.py, root install in `/usr/local/bin`:
-```
-# ./setup.py install
-```
-installation using setup.py, user install in `~/.local/bin`:
-```
-# ./setup.py install --prefix=~/.local
-```
+
 ## Usage
 ```
 # pgtree -h
@@ -58,6 +53,7 @@ installation using setup.py, user install in `~/.local/bin`:
     -w : tty wrap text : y/yes or n/no (default y)
     -W : watch and follow process tree every 2s
     -a : use ascii characters
+    -T : display threads (ps -T)
     -O <psfield>[,psfield,...] : display multiple <psfield> instead of 'stime' in output
                    <psfield> must be valid with ps -o <psfield> command
 

pgtree/pgtree.py

@@ -2,9 +2,21 @@
 # coding: utf-8
 # pylint: disable=C0114,C0413,R0902,C0209
 # determine available python executable
+# determine ps -o options
 _=''''
 #[ "$1" = -W ] && shift && exec watch -x -c -- "$0" -C y "$@"
-export PGT_PGREP=$(type -p pgrep)
+export LANG=en_US.UTF-8 PYTHONUTF8=1 PYTHONIOENCODING=utf8
+PGT_PGREP=$(type -p pgrep)
+ps -p $$ -o ucomm >/dev/null 2>&1 && PGT_COMM=ucomm
+[ ! "$PGT_COMM" ] && ps -p $$ -o comm >/dev/null 2>&1 && PGT_COMM=comm
+[ "$PGT_COMM" ] && {
+    ps -p $$ -o stime >/dev/null 2>&1 && PGT_STIME=stime
+    [ ! "$PGT_STIME" ] && ps -p $$ -o start >/dev/null 2>&1 && PGT_STIME=start
+    [ ! "$PGT_STIME" ] && PGT_STIME=time
+}
+# busybox no -p option
+[ ! "$PGT_COMM" ] && ! ps -p $$ >/dev/null 2>&1 && PGT_COMM=comm && PGT_STIME=time
+export PGT_COMM PGT_STIME PGT_PGREP
 python=$(type -p python || type -p python3 || type -p python2)
 [ "$python" ] && exec $python "$0" "$@"
 echo "ERROR: cannot find python interpreter" >&2
@@ -17,7 +29,7 @@
 hierarchy of matching processes (parents and children)
 should work on any Unix supporting commands :
 # pgrep
-# ps -e -o pid,ppid,comm,args
+# ps ax -o pid,ppid,comm,args
 (RedHat/CentOS/Fedora/Ubuntu/Suse/Solaris...)
 Compatible python 2 / 3
 
@@ -45,23 +57,30 @@
 
 import sys
 import os
-import getopt
 import platform
+import getopt
 import re
-import time
-
-# pylint: disable=E0602
-# pylint: disable=E1101
-if sys.version_info < (3, 0):
-    reload(sys)
-    sys.setdefaultencoding('utf8')
+try:
+    import time
+except ImportError:
+    pass
+
+# impossible detection using ps for AIX/MacOS
+# stime is not start time of process
+system = platform.system()
+PS_OPTION = 'ax'
+if system in ['AIX', 'Darwin']:
+    os.environ['PGT_STIME'] = 'start'
+elif system == 'SunOS': # ps ax -o not supported
+    PS_OPTION = '-e'
+    os.environ['PGT_COMM'] = 'fname' # comm header width not respected
 
 def runcmd(cmd):
     """run command"""
-    pipe = os.popen('"' + '" "'.join(cmd) + '"', 'r')
+    pipe = os.popen(cmd, 'r')
     std_out = pipe.read()
-    pipe.close()
-    return std_out.rstrip('\n')
+    res = pipe.close()
+    return res, std_out.rstrip('\n')
 
 def ask(prompt):
     """input text"""
@@ -119,7 +138,7 @@ class Proctree:
 
     # pylint: disable=R0913
     def __init__(self, use_uid=False, use_ascii=False, use_color=False,
-                 pid_zero=True, opt_fields=None):
+                 pid_zero=True, opt_fields=None, threads=False):
         """constructor"""
         self.pids = []
         self.ps_info = {}        # ps command info stored
@@ -128,52 +147,84 @@ def __init__(self, use_uid=False, use_ascii=False, use_color=False,
         self.pids_tree = {}
         self.top_parents = []
         self.treedisp = Treedisplay(use_ascii, use_color)
-        self.ps_fields = self.get_fields(opt_fields, use_uid)
+        self.ps_fields = self.get_fields(opt_fields, use_uid, threads)
         self.get_psinfo(pid_zero)
 
-    def get_fields(self, opt_fields=None, use_uid=False):
+    def get_fields(self, opt_fields=None, use_uid=False, threads=False):
         """ Get ps fields from OS / optionnal fields """
-        osname = platform.system()
-        if not opt_fields:
-            if osname in ['AIX', 'Darwin']:
-                opt_fields = ['start']
-            else:
-                opt_fields = ['stime']
         if use_uid:
             user = 'uid'
         else:
             user = 'user'
-        if osname == 'SunOS':
-            comm = 'comm'
+        if threads:
+            pid = 'spid'
         else:
-            comm = 'ucomm'
-
-        return ['pid', 'ppid', user, comm] + opt_fields
+            pid = 'pid'
+        if not opt_fields or not os.environ.get('PGT_COMM'):
+            opt_fields = [os.environ.get('PGT_STIME') or 'stime']
+
+        return [pid, 'ppid', user, os.environ.get('PGT_COMM') or 'ucomm'] + opt_fields
+
+    def run_ps(self, widths):
+        """
+            run ps command detected setting columns widths
+            guess columns for ps command not supporting -o (mingw/msys2)
+        """
+        if os.environ.get('PGT_COMM'):
+            ps_cmd = 'ps ' + PS_OPTION + ' ' + ' '.join(
+                    ['-o '+ o +'='+ widths[i]*'-' for i,o in enumerate(self.ps_fields)]
+                ) + ' -o args'
+            err, ps_out = runcmd(ps_cmd)
+            if err:
+                print('Error: executing ps ' + PS_OPTION + ' -o ' + ",".join(self.ps_fields))
+                sys.exit(1)
+            return ps_out.splitlines()
+        _, out = runcmd('ps aux') # try to use header to guess columns
+        out = out.splitlines()
+        if not 'PPID' in out[0]:
+            _, out = runcmd('ps -ef')
+            out = out.splitlines()
+        ps_out = []
+        fields = {}
+        for i,field in enumerate(out[0].strip().lower().split()):
+            field = re.sub("command|cmd", "args", field)
+            field = re.sub("uid", "user", field)
+            fields[field] = i
+        if not 'ppid' in fields:
+            print("Error: command 'ps aux' does not provides PPID")
+            sys.exit(1)
+        fields["ucomm"] = len(fields)
+        for line in out:
+            ps_info = line.strip().split(None, len(fields)-2)
+            if "stime" in fields:
+                if ps_info[fields["stime"]] in ["Jan","Feb","Mar","Apr","May","Jun",
+                                           "Jui","Aug","Sep","Oct","Nov","Dec"]:
+                    ps_info = line.strip().split(None, len(fields)-1)
+                    ps_info[fields["stime"]] += ps_info.pop(fields["stime"]+1)
+            ps_info.append(os.path.basename(ps_info[fields["args"]].split()[0]))
+            ps_out.append(' '.join(
+                [('%-'+ str(widths[i]) +'s') % ps_info[fields[opt]]
+                 for i,opt in enumerate(self.ps_fields)] + [ps_info[fields["args"]]]
+            ))
+        return ps_out
 
     def get_psinfo(self, pid_zero):
         """parse unix ps command"""
         widths = [30, 30, 30, 130] + [50 for i in self.ps_fields[4:]]
-        ps_cmd = 'ps -e ' + ' '.join(
-                    ['-o '+ o +'='+ widths[i]*'-' for i,o in enumerate(self.ps_fields)]
-                ) + ' -o args'
-        # print(ps_cmd)
-        ps_out = runcmd(ps_cmd.split(' ')).split('\n')
+        ps_out = self.run_ps(widths)
         pid_z = ["0", "0"] + self.ps_fields[2:]
         ps_out[0] = ' '.join(
                 [('%-'+ str(widths[i]) +'s') % opt for i,opt in enumerate(pid_z)] + ['args']
         )
         ps_opts = ['pid', 'ppid', 'user', 'comm'] + self.ps_fields[4:]
-        # print(ps_out[0])
         for line in ps_out:
-            # print(line)
             infos = {}
             col = 0
             for i,field in enumerate(ps_opts):
                 infos[field] = line[col:col+widths[i]].strip()
                 col = col + widths[i] + 1
             infos['args'] = line[col:len(line)]
             infos['comm'] = os.path.basename(infos['comm'])
-            # print(infos)
             pid = infos['pid']
             ppid = infos['ppid']
             if pid == str(os.getpid()):
@@ -185,6 +236,8 @@ def get_psinfo(self, pid_zero):
                 self.children[ppid] = []
             self.children[ppid].append(pid)
             self.ps_info[pid] = infos
+        if not self.ps_info.get('1'):
+            self.ps_info['1'] = self.ps_info['0']
         if not pid_zero:
             del self.ps_info['0']
             del self.children['0']
@@ -193,7 +246,7 @@ def pgrep(self, argv):
         """mini built-in pgrep if pgrep command not available
            [-f] [-x] [-i] [-u <user>] [pattern]"""
         if "PGT_PGREP" not in os.environ or os.environ["PGT_PGREP"]:
-            pgrep = runcmd(['pgrep'] + argv)
+            _, pgrep = runcmd('pgrep ' +' '.join(argv))
             return pgrep.split("\n")
 
         try:
@@ -232,6 +285,7 @@ def pgrep(self, argv):
 
     def get_parents(self):
         """get parents list of pids"""
+        last_ppid = None
         for pid in self.pids:
             if pid not in self.ps_info:
                 continue
@@ -361,7 +415,8 @@ def pgtree(options, psfields, pgrep_args):
                      use_ascii='-a' in options,
                      use_color=colored(options['-C']),
                      pid_zero='-1' not in options,
-                     opt_fields=psfields)
+                     opt_fields=psfields,
+                     threads='-T' in options)
 
     found = None
     if '-p' in options:
@@ -389,6 +444,7 @@ def watch_pgtree(options, psfields, pgrep_args, sig):
 
 def main(argv):
     """pgtree command line"""
+    global PS_OPTION
     usage = """
     usage: pgtree.py [-W] [-RIya] [-C <when>] [-O <psfield>] [-c|-k|-K] [-1|-p <pid1>,...|<pgrep args>]
 
@@ -403,6 +459,7 @@ def main(argv):
     -w : tty wrap text : y/yes or n/no (default y)
     -W : watch and follow process tree every 2s
     -a : use ascii characters
+    -T : display threads (ps -T)
     -O <psfield>[,psfield,...] : display multiple <psfield> instead of 'stime' in output
                    <psfield> must be valid with ps -o <psfield> command
 
@@ -422,7 +479,7 @@ def main(argv):
         argv = os.environ["PGTREE"].split(' ') + argv
     try:
         opts, args = getopt.getopt(argv,
-                                   "W1IRckKfxvinoyap:u:U:g:G:P:s:t:F:O:C:w:",
+                                   "W1IRckKfxvinoyaTp:u:U:g:G:P:s:t:F:O:C:w:",
                                    ["ns=", "nslist="])
     except getopt.GetoptError:
         print(usage)
@@ -444,6 +501,8 @@ def main(argv):
             psfields = arg.split(',')
         elif opt == "-R":
             os.environ["PGT_PGREP"] = ""
+        elif opt == "-T":
+            PS_OPTION += " -T"
         elif opt in ("-f", "-x", "-v", "-i", "-n", "-o"):
             pgrep_args.append(opt)
         elif opt in ("-u", "-U", "-g", "-G", "-P", "-s", "-t", "-F", "--ns", "--nslist"):

tests/test_pgtree.py

@@ -6,6 +6,8 @@
 sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
 import pgtree
 #from unittest.mock import MagicMock, Mock, patch
+os.environ['PGT_COMM'] = 'ucomm'
+os.environ['PGT_STIME'] = 'stime'
 
 class ProctreeTest(unittest.TestCase):
     """tests for pgtree"""
@@ -22,7 +24,7 @@ def test_tree1(self, mock_runcmd, mock_kill):
         ps_out += f'{"30":>30} {"10":>30} {"joknarf":<30} {"top":<130} {"10:10":<50} /bin/top\n'
         ps_out += f'{"40":>30} {"1":>30} {"root":<30} {"bash":<130} {"11:01":<50} -bash'
         print(ps_out)
-        mock_runcmd.return_value = ps_out
+        mock_runcmd.return_value = 0, ps_out
         mock_kill.return_value = True
         ptree = pgtree.Proctree()
 
@@ -140,7 +142,7 @@ def test_main6(self):
     def test_main7(self):
         """test"""
         print('main7 ========')
-        pgtree.main(['-O', '%cpu', 'bash'])
+        pgtree.main(['-C', 'y', '-O', '%cpu', 'init'])
 
     def test_ospgrep(self):
         """pgrep os"""
@@ -165,6 +167,17 @@ def test_watch(self, mock_sleep):
         mock_sleep.return_value = True
         pgtree.main(['-W', 'bash'])
 
+    @patch.dict(os.environ, {"PGT_COMM": "", "PGT_STIME": ""})
+    def test_simpleps(self):
+        pgtree.main([])
+
+    def test_psfail(self):
+        """test"""
+        print('psfail ========')
+        try:
+            pgtree.main(['-O abcd'])
+        except SystemExit:
+            pass
 
-if __name__ == "__main__":
-    unittest.main(failfast=True)
+    def test_threads(self):
+        pgtree.main(["-T"])