From 40b49768d6c2776d9303d1bc6eb0f6021dec5479 Mon Sep 17 00:00:00 2001
From: Matt Mackall <mpm@selenic.com>
Date: Sat, 4 Apr 2009 16:33:00 -0500
Subject: [PATCH] Abstract proc interface

---
 smem | 81 +++++++++++++++++++++++++++++++++++++-----------------------
 1 file changed, 50 insertions(+), 31 deletions(-)

diff --git a/smem b/smem
index 82afab3..44cad1b 100755
--- a/smem
+++ b/smem
@@ -1,6 +1,41 @@
 #!/usr/bin/python
 import re, os, sys, pwd, grp, optparse, errno
 
+_ucache = {}
+_gcache = {}
+
+class procdata(object):
+    def pids(self):
+        '''get a list of processes'''
+        return [int(e) for e in os.listdir("/proc")
+                if e.isdigit() and not iskernel(e)]
+    def _read(self, f):
+        return file('/proc/' + f).read()
+    def _readlines(self, f):
+        return self._read(f).splitlines(True)
+    def mapdata(self, pid):
+        return self._readlines('%s/smaps' % pid)
+    def memdata(self):
+        return self._readlines('meminfo')
+    def pidname(self, pid):
+        l = self._read('%d/stat' % pid)
+        return l[l.find('(') + 1: l.find(')')]
+    def pidcmd(self, pid):
+        c = self._read('%s/cmdline' % pid)[:-1]
+        return c.replace('\0', ' ')
+    def piduser(self, pid):
+        return os.stat('/proc/%d/cmdline' % pid).st_uid
+    def pidgroup(self, pid):
+        return os.stat('/proc/%d/cmdline' % pid).st_gid
+    def username(self, uid):
+        if uid not in _ucache:
+            _ucache[uid] = pwd.getpwuid(uid)[0]
+        return _ucache[uid]
+    def groupname(self, gid):
+        if gid not in _gcache:
+            _gcache[gid] = pwd.getgrgid(gid)[0]
+        return _gcache[gid]
+
 _totalmem = 0
 def totalmem():
     global _totalmem
@@ -11,14 +46,10 @@ def totalmem():
             _totalmem = memory()['memtotal']
     return _totalmem
 
-def pids():
-    '''get a list of processes'''
-    return [int(e) for e in os.listdir("/proc") if e.isdigit() and not iskernel(e)]
-
 def pidmaps(pid):
     maps = {}
     start = None
-    for l in file('/proc/%s/smaps' % pid):
+    for l in src.mapdata(pid):
     	f = l.split()
 	if f[-1] == 'kB':
             maps[start][f[0][:-1].lower()] = int(f[1])
@@ -49,26 +80,12 @@ def sortmaps(totals, key):
     return [pid for pid,key in l]
 
 def iskernel(pid):
-    return pidcmd(pid) == ""
-
-def pidname(pid):
-    l = file('/proc/%d/stat' % pid).read()
-    return l[l.find('(') + 1: l.find(')')]
-
-def pidcmd(pid):
-    c = file('/proc/%s/cmdline' % pid).read()[:-1]
-    return c.replace('\0', ' ')
-
-def piduser(pid):
-    return os.stat('/proc/%d/cmdline' % pid).st_uid
-
-def pidgroup(pid):
-    return os.stat('/proc/%d/cmdline' % pid).st_gid
+    return src.pidcmd(pid) == ""
 
 def memory():
     t = {}
     f = re.compile('(\\S+):\\s+(\\d+) kB')
-    for l in file('/proc/meminfo'):
+    for l in src.memdata():
         m = f.match(l)
         if m:
             t[m.group(1).lower()] = int(m.group(2)) * 1024
@@ -99,7 +116,7 @@ def username(uid):
     return _ucache[uid]
 
 def pidusername(pid):
-    return username(piduser(pid))
+    return username(src.piduser(pid))
 
 _gcache = {}
 def groupname(gid):
@@ -146,7 +163,7 @@ def pidtotals(pid):
 def processtotals(pids):
     totals = {}
     for pid in pids:
-        if (filter(options.processfilter, pid, pidname, pidcmd) or
+        if (filter(options.processfilter, pid, src.pidname, src.pidcmd) or
                 filter(options.userfilter, pid, pidusername)):
             continue
         try:
@@ -158,7 +175,7 @@ def processtotals(pids):
     return totals
 
 def showpids():
-    p = pids()
+    p = src.pids()
     pt = processtotals(p)
 
     def showuser(p):
@@ -169,8 +186,8 @@ def showpids():
     fields = dict(
         pid=('PID', lambda n: n, '% 5s', lambda x: len(p)),
         user=('User', showuser, '%-8s', lambda x: len(dict.fromkeys(x))),
-        name=('Name', pidname, '%-24.24s', None),
-        command=('Command', pidcmd, '%-27.27s', None),
+        name=('Name', src.pidname, '%-24.24s', None),
+        command=('Command', src.pidcmd, '%-27.27s', None),
         maps=('Maps',lambda n: pt[n]['maps'], '% 5s', sum),
         swap=('Swap',lambda n: pt[n]['swap'], '% 8a', sum),
         uss=('USS', lambda n: pt[n]['uss'], '% 8a', sum),
@@ -189,7 +206,7 @@ def showpids():
 def maptotals(pids):
     totals = {}
     for pid in pids:
-        if (filter(options.processfilter, pid, pidname, pidcmd) or
+        if (filter(options.processfilter, pid, src.pidname, src.pidcmd) or
                 filter(options.userfilter, pid, pidusername)):
             continue
         try:
@@ -212,7 +229,7 @@ def maptotals(pids):
     return totals
 
 def showmaps():
-    p = pids()
+    p = src.pids()
     pt = maptotals(p)
 
     fields = dict(
@@ -238,7 +255,7 @@ def showmaps():
 def usertotals(pids):
     totals = {}
     for pid in pids:
-        if (filter(options.processfilter, pid, pidname, pidcmd) or
+        if (filter(options.processfilter, pid, src.pidname, src.pidcmd) or
                 filter(options.userfilter, pid, pidusername)):
             continue
         try:
@@ -247,7 +264,7 @@ def usertotals(pids):
                 continue
         except:
             raise
-        user = piduser(pid)
+        user = src.piduser(pid)
         if user not in totals:
             t = dict(size=0, rss=0, pss=0, shared_clean=0,
                      shared_dirty=0, private_clean=0, count=0,
@@ -264,7 +281,7 @@ def usertotals(pids):
     return totals
 
 def showusers():
-    p = pids()
+    p = src.pids()
     pt = usertotals(p)
 
     def showuser(u):
@@ -364,6 +381,8 @@ defaults = {}
 parser.set_defaults(**defaults)
 (options, args) = parser.parse_args()
 
+src = procdata()
+
 try:
     if options.mappings:
         showmaps()