Add code to measure memory usage on windows

src/calibre/utils/mem.py

@@ -5,53 +5,91 @@ __license__   = 'GPL v3'
 __copyright__ = '2010, Kovid Goyal <kovid@kovidgoyal.net>'
 __docformat__ = 'restructuredtext en'
 
-import gc
+import gc, os
 
-## {{{ http://code.activestate.com/recipes/286222/ (r1)
+from calibre.constants import iswindows, islinux
-import os
 
-_proc_status = '/proc/%d/status' % os.getpid()
+if islinux:
+    ## {{{ http://code.activestate.com/recipes/286222/ (r1)
 
-_scale = {'kB': 1024.0, 'mB': 1024.0*1024.0,
+    _proc_status = '/proc/%d/status' % os.getpid()
-          'KB': 1024.0, 'MB': 1024.0*1024.0}
 
-def _VmB(VmKey):
+    _scale = {'kB': 1024.0, 'mB': 1024.0*1024.0,
-    '''Private.
+            'KB': 1024.0, 'MB': 1024.0*1024.0}
-    '''
+
-    global _proc_status, _scale
+    def _VmB(VmKey):
-     # get pseudo file  /proc/<pid>/status
+        '''Private.
-    try:
+        '''
-        t = open(_proc_status)
+        global _proc_status, _scale
-        v = t.read()
+        # get pseudo file  /proc/<pid>/status
-        t.close()
+        try:
-    except:
+            t = open(_proc_status)
-        return 0.0  # non-Linux?
+            v = t.read()
-     # get VmKey line e.g. 'VmRSS:  9999  kB\n ...'
+            t.close()
-    i = v.index(VmKey)
+        except:
-    v = v[i:].split(None, 3)  # whitespace
+            return 0.0  # non-Linux?
-    if len(v) < 3:
+        # get VmKey line e.g. 'VmRSS:  9999  kB\n ...'
-        return 0.0  # invalid format?
+        i = v.index(VmKey)
-     # convert Vm value to bytes
+        v = v[i:].split(None, 3)  # whitespace
-    return float(v[1]) * _scale[v[2]]
+        if len(v) < 3:
+            return 0.0  # invalid format?
+        # convert Vm value to bytes
+        return float(v[1]) * _scale[v[2]]
 
 
-def memory(since=0.0):
+    def linux_memory(since=0.0):
-    '''Return memory usage in bytes.
+        '''Return memory usage in bytes.
-    '''
+        '''
-    return _VmB('VmSize:') - since
+        return _VmB('VmSize:') - since
 
 
-def resident(since=0.0):
+    def resident(since=0.0):
-    '''Return resident memory usage in bytes.
+        '''Return resident memory usage in bytes.
-    '''
+        '''
-    return _VmB('VmRSS:') - since
+        return _VmB('VmRSS:') - since
 
 
-def stacksize(since=0.0):
+    def stacksize(since=0.0):
-    '''Return stack size in bytes.
+        '''Return stack size in bytes.
-    '''
+        '''
-    return _VmB('VmStk:') - since
+        return _VmB('VmStk:') - since
-## end of http://code.activestate.com/recipes/286222/ }}}
+    ## end of http://code.activestate.com/recipes/286222/ }}}
+    memory = linux_memory
+elif iswindows:
+    import win32process
+    import win32con
+    import win32api
+
+    # See http://msdn.microsoft.com/en-us/library/ms684877.aspx
+    # for details on the info returned by get_meminfo
+
+    def get_handle(pid):
+        return win32api.OpenProcess(win32con.PROCESS_QUERY_INFORMATION, 0,
+                pid)
+
+    def listprocesses(self):
+        for process in win32process.EnumProcesses():
+            try:
+                han = get_handle(process)
+                procmeminfo = meminfo(han)
+                procmemusage = procmeminfo["WorkingSetSize"]
+                yield process, procmemusage
+            except:
+                pass
+
+    def get_meminfo(pid):
+        han = win32api.OpenProcess(win32con.PROCESS_QUERY_INFORMATION, 0,
+                pid)
+        return meminfo(han)
+
+    def meminfo(handle):
+        return win32process.GetProcessMemoryInfo(handle)
+
+    def win_memory(since=0.0):
+        info = meminfo(get_handle(os.getpid()))
+        return info['WorkingSetSize'] - since
+
+    memory = win_memory
 
 
 def gc_histogram():