Issue 142: implement process nice/priority on Windows.

6 files changed, 142 insertions, 34 deletions

diff --git a/psutil/__init__.py b/psutil/__init__.py
index 577b75ef..438de7dc 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -161,6 +161,16 @@ class Process(object):
         """The command line process has been called with."""
         return self._platform_impl.get_process_cmdline()
 
+    @property
+    def nice(self):
+        """Get or set process niceness (priority)."""
+        return self._platform_impl.get_process_nice()
+
+    @nice.setter
+    def nice(self, value):
+        # invoked on "p.nice = n", change process niceness
+        return self._platform_impl.set_process_nice(value)
+
     if os.name == 'posix':
 
         @property
@@ -177,16 +187,6 @@ class Process(object):
             """
             return self._platform_impl.get_process_gids()
 
-        @property
-        def nice(self):
-            """Get or set process niceness."""
-            return self._platform_impl.get_process_nice()
-
-        @nice.setter
-        def nice(self, value):
-            # invoked on "p.nice = n", change process niceness
-            return self._platform_impl.set_process_nice(value)
-
     @property
     def uid(self):
         """The real user id of the current process."""
diff --git a/psutil/_psmswindows.py b/psutil/_psmswindows.py
index 8c84c933..a3a04708 100644
--- a/psutil/_psmswindows.py
+++ b/psutil/_psmswindows.py
@@ -19,17 +19,24 @@ except ImportError:
 import _psutil_mswindows
 from psutil.error import AccessDenied, NoSuchProcess
 
-
-
 # --- module level constants (gets pushed up to psutil module)
 
 NUM_CPUS = _psutil_mswindows.get_num_cpus()
 TOTAL_PHYMEM = _psutil_mswindows.get_total_phymem()
 BOOT_TIME = _psutil_mswindows.get_system_uptime()
 _WIN2000 = platform.win32_ver()[0] == '2000'
-
 ERROR_ACCESS_DENIED = 5
 
+# process priority constants:
+# http://msdn.microsoft.com/en-us/library/ms686219(v=vs.85).aspx
+from _psutil_mswindows import (ABOVE_NORMAL_PRIORITY_CLASS,
+                               BELOW_NORMAL_PRIORITY_CLASS,
+                               HIGH_PRIORITY_CLASS,
+                               IDLE_PRIORITY_CLASS,
+                               NORMAL_PRIORITY_CLASS,
+                               REALTIME_PRIORITY_CLASS)
+
+
 # --- public functions
 
 def avail_phymem():
@@ -226,4 +233,13 @@ class WindowsProcess(object):
         retlist = _psutil_mswindows.get_process_connections(self.pid)
         return [self._connection_ntuple(*conn) for conn in retlist]
 
+    @wrap_exceptions
+    def get_process_nice(self):
+        return _psutil_mswindows.get_process_priority(self.pid)
+
+    @wrap_exceptions
+    def set_process_nice(self, value):
+        return _psutil_mswindows.set_process_priority(self.pid, value)
+
+
 PlatformProcess = WindowsProcess
diff --git a/psutil/_psutil_mswindows.c b/psutil/_psutil_mswindows.c
index 42bf9dfd..3c7d2542 100644
--- a/psutil/_psutil_mswindows.c
+++ b/psutil/_psutil_mswindows.c
@@ -62,6 +62,10 @@ PsutilMethods[] =
         "Return process threads information as a list of tuple"},
     {"process_wait", process_wait, METH_VARARGS,
         "Wait for process to terminate and return its exit code."},
+    {"get_process_priority", get_process_priority, METH_VARARGS,
+        "Return process priority."},
+    {"set_process_priority", set_process_priority, METH_VARARGS,
+        "Set process priority."},
 
     // --- system-related functions
 
@@ -150,6 +154,20 @@ struct module_state {
         INITERROR;
     }
 
+    // Public constants
+    // http://msdn.microsoft.com/en-us/library/ms683211(v=vs.85).aspx
+    PyModule_AddIntConstant(module, "ABOVE_NORMAL_PRIORITY_CLASS",
+                                     ABOVE_NORMAL_PRIORITY_CLASS);
+    PyModule_AddIntConstant(module, "BELOW_NORMAL_PRIORITY_CLASS",
+                                     BELOW_NORMAL_PRIORITY_CLASS);
+    PyModule_AddIntConstant(module, "HIGH_PRIORITY_CLASS",
+                                     HIGH_PRIORITY_CLASS);
+    PyModule_AddIntConstant(module, "IDLE_PRIORITY_CLASS",
+                                     IDLE_PRIORITY_CLASS);
+    PyModule_AddIntConstant(module, "NORMAL_PRIORITY_CLASS",
+                                     NORMAL_PRIORITY_CLASS);
+    PyModule_AddIntConstant(module, "REALTIME_PRIORITY_CLASS",
+                                     REALTIME_PRIORITY_CLASS);
     SetSeDebug();
 
 #if PY_MAJOR_VERSION >= 3
@@ -1810,3 +1828,59 @@ get_process_connections(PyObject* self, PyObject* args)
     return connectionsList;
 }
 
+
+/*
+ * Get process priority as a Python integer.
+ */
+static PyObject*
+get_process_priority(PyObject* self, PyObject* args)
+{
+    long pid;
+    DWORD priority;
+    HANDLE hProcess;
+    if (! PyArg_ParseTuple(args, "l", &pid)) {
+        return NULL;
+    }
+
+    hProcess = handle_from_pid(pid);
+    if (hProcess == NULL) {
+        return NULL;
+    }
+
+    priority = GetPriorityClass(hProcess);
+    if (priority == 0) {
+        PyErr_SetFromWindowsErr(0);
+        return NULL;
+    }
+    return Py_BuildValue("i", priority);
+}
+
+
+/*
+ * Set process priority.
+ */
+static PyObject*
+set_process_priority(PyObject* self, PyObject* args)
+{
+    long pid;
+    int priority;
+    int retval;
+    HANDLE hProcess;
+    DWORD dwDesiredAccess = PROCESS_QUERY_INFORMATION | PROCESS_SET_INFORMATION;
+    if (! PyArg_ParseTuple(args, "li", &pid, &priority)) {
+        return NULL;
+    }
+
+    hProcess = handle_from_pid_waccess(pid, dwDesiredAccess);
+    if (hProcess == NULL) {
+        return NULL;
+    }
+
+    retval = SetPriorityClass(hProcess, priority);
+    if (retval == 0) {
+        PyErr_SetFromWindowsErr(0);
+        return NULL;
+    }
+    Py_INCREF(Py_None);
+    return Py_None;
+}
diff --git a/psutil/_psutil_mswindows.h b/psutil/_psutil_mswindows.h
index b6e0457c..fdd855aa 100644
--- a/psutil/_psutil_mswindows.h
+++ b/psutil/_psutil_mswindows.h
@@ -26,6 +26,8 @@ static PyObject* get_process_connections(PyObject* self, PyObject* args);
 static PyObject* get_process_num_threads(PyObject* self, PyObject* args);
 static PyObject* get_process_threads(PyObject* self, PyObject* args);
 static PyObject* process_wait(PyObject* self, PyObject* args);
+static PyObject* get_process_priority(PyObject* self, PyObject* args);
+static PyObject* set_process_priority(PyObject* self, PyObject* args);
 
 // --- system-related functions
 
diff --git a/setup.py b/setup.py
index b6bccf14..c87d7dd9 100644
--- a/setup.py
+++ b/setup.py
@@ -29,17 +29,17 @@ if sys.platform.lower().startswith("win"):
         maj,min = sys.getwindowsversion()[0:2]
         return '0x0%s' % ((maj * 100) + min)
 
-    extensions = Extension('_psutil_mswindows',
-                           sources=['psutil/_psutil_mswindows.c',
-                                    'psutil/_psutil_common.c',
-                                    'psutil/arch/mswindows/process_info.c',
-                                    'psutil/arch/mswindows/process_handles.c',
-                                    'psutil/arch/mswindows/security.c'],
-                           define_macros=[('_WIN32_WINNT', get_winver()),
-                                          ('_AVAIL_WINVER_', get_winver())],
-                           libraries=["psapi", "kernel32", "advapi32", "shell32",
-                                      "netapi32"]
-                           )
+    extensions = [Extension('_psutil_mswindows',
+                            sources=['psutil/_psutil_mswindows.c',
+                                     'psutil/_psutil_common.c',
+                                     'psutil/arch/mswindows/process_info.c',
+                                     'psutil/arch/mswindows/process_handles.c',
+                                     'psutil/arch/mswindows/security.c'],
+                            define_macros=[('_WIN32_WINNT', get_winver()),
+                                           ('_AVAIL_WINVER_', get_winver())],
+                            libraries=["psapi", "kernel32", "advapi32", "shell32",
+                                       "netapi32"]
+                            )]
 # OS X
 elif sys.platform.lower().startswith("darwin"):
     extensions = [Extension('_psutil_osx',
diff --git a/test/test_psutil.py b/test/test_psutil.py
index efc9a9f0..d5425264 100644
--- a/test/test_psutil.py
+++ b/test/test_psutil.py
@@ -600,6 +600,7 @@ class TestCase(unittest.TestCase):
 
         def test_nice(self):
             p = psutil.Process(os.getpid())
+            self.assertRaises(TypeError, setattr, p, "nice", "str")
             first_nice = p.nice
             try:
                 p.nice = 1
@@ -609,6 +610,20 @@ class TestCase(unittest.TestCase):
             finally:
                 p.nice = first_nice
 
+    if os.name == 'nt':
+
+        def test_nice(self):
+            p = psutil.Process(os.getpid())
+            self.assertRaises(TypeError, setattr, p, "nice", "str")
+            try:
+                self.assertEqual(p.nice, psutil.NORMAL_PRIORITY_CLASS)
+                p.nice = psutil.HIGH_PRIORITY_CLASS
+                self.assertEqual(p.nice, psutil.HIGH_PRIORITY_CLASS)
+                p.nice = psutil.NORMAL_PRIORITY_CLASS
+                self.assertEqual(p.nice, psutil.NORMAL_PRIORITY_CLASS)
+            finally:
+                p.nice = psutil.NORMAL_PRIORITY_CLASS
+
     def test_username(self):
         sproc = get_test_subprocess()
         p = psutil.Process(sproc.pid)
@@ -1009,7 +1024,10 @@ class TestCase(unittest.TestCase):
         if hasattr(p, 'nice'):
             self.assertRaises(psutil.NoSuchProcess, p.nice)
             try:
-                p.nice = 1
+                if os.name == 'posix':
+                    p.nice = 1
+                else:
+                    p.nice = psutil.NORMAL_PRIORITY_CLASS
             except psutil.NoSuchProcess:
                 pass
             else:
@@ -1147,15 +1165,13 @@ if hasattr(os, 'getuid'):
             # DeprecationWarning is only raised once
             pass
 
-        if os.name == 'posix':
-
-            def test_nice(self):
-                try:
-                    psutil.Process(os.getpid()).nice = -1
-                except psutil.AccessDenied:
-                    pass
-                else:
-                    self.fail("exception not raised")
+        def test_nice(self):
+            try:
+                psutil.Process(os.getpid()).nice = -1
+            except psutil.AccessDenied:
+                pass
+            else:
+                self.fail("exception not raised")
 
         # overridden tests known to raise AccessDenied when run
         # as limited user on different platforms