#1473: implement ionice() on OSX

6 files changed, 167 insertions, 21 deletions

diff --git a/docs/index.rst b/docs/index.rst
index 6ff1e22c..65646315 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -1230,13 +1230,21 @@ Process class
 
   .. method:: ionice(ioclass=None, value=None)
 
-    Get or set process I/O niceness (priority). On Linux *ioclass* is one of the
+    Get or set process I/O niceness (priority).
+
+    On Linux *ioclass* is one of the
     :data:`psutil.IOPRIO_CLASS_*<psutil.IOPRIO_CLASS_NONE>` constants.
     *value* is a number which goes from  ``0`` to ``7``. The higher the value,
-    the lower the I/O priority of the process. On Windows only *ioclass* is
-    used and it can be set to ``2`` (normal), ``1`` (low) or ``0`` (very low).
-    The example below sets IDLE priority class for the current process,
-    meaning it will only get I/O time when no other process needs the disk:
+    the lower the I/O priority of the process.
+    On Linux this returns a ``(ioclass, value))`` tuple, on all other platforms
+    it returns an integer.
+    On macOS only *ioclass* is used and it must be one of the
+    :data:`psutil.IOPOL_*<psutil.IOPOL_DEFAULT>` constants.
+    On Windows only *ioclass* is used and can be set to ``2`` (normal),
+    ``1`` (low) or ``0`` (very low).
+    The example below sets IDLE priority class for the current process on
+    Linux, meaning it will only get I/O time when no other process needs the
+    disk:
 
       >>> import psutil
       >>> p = psutil.Process()
@@ -1245,11 +1253,9 @@ Process class
       pionice(ioclass=<IOPriority.IOPRIO_CLASS_IDLE: 3>, value=0)
       >>>
 
-    On Windows only *ioclass* is used and it can be set to ``2`` (normal),
-    ``1`` (low) or ``0`` (very low). Also it returns an integer instead of a
-    named tuple.
+    Availability: Linux, macOS, Windows Vista+
 
-    Availability: Linux and Windows > Vista
+    .. versionchanged:: 5.6.2 added macOS support
 
     .. versionchanged::
       3.0.0 on Python >= 3.4 the returned ``ioclass`` constant is an
@@ -2162,6 +2168,21 @@ Constants
     `enums <https://docs.python.org/3/library/enum.html#module-enum>`__
     instead of a plain integer.
 
+.. _const-iopol:
+.. data:: IOPOL_DEFAULT
+.. data:: IOPOL_IMPORTANT
+.. data:: IOPOL_PASSIVE
+.. data:: IOPOL_THROTTLE
+.. data:: IOPOL_UTILITY
+.. data:: IOPOL_STANDARD
+
+  A set of integers representing the I/O priority of a process on macOS. They
+  can be used in conjunction with :meth:`psutil.Process.ionice()` to get or set
+  process I/O priority.  For further information refer to
+  `ionice <http://www.manpagez.com/man/3/setiopolicy_np/>`__ manual.
+
+  Availability: macOS
+
 .. _const-rlimit:
 .. data:: RLIM_INFINITY
 .. data:: RLIMIT_AS
diff --git a/psutil/__init__.py b/psutil/__init__.py
index ab2ed349..18fbe71b 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -149,6 +149,12 @@ elif WINDOWS:
 
 elif MACOS:
     from . import _psosx as _psplatform
+    from ._psosx import IOPOL_DEFAULT  # NOQA
+    from ._psosx import IOPOL_IMPORTANT  # NOQA
+    from ._psosx import IOPOL_PASSIVE  # NOQA
+    from ._psosx import IOPOL_THROTTLE  # NOQA
+    from ._psosx import IOPOL_UTILITY  # NOQA
+    from ._psosx import IOPOL_STANDARD  # NOQA
 
 elif BSD:
     from . import _psbsd as _psplatform
@@ -871,7 +877,7 @@ class Process(object):
             """
             return self._proc.io_counters()
 
-    # Linux and Windows >= Vista only
+    # Linux, macOS and Windows >= Vista only
     if hasattr(_psplatform.Process, "ionice_get"):
 
         def ionice(self, ioclass=None, value=None):
@@ -1475,6 +1481,8 @@ _as_dict_attrnames = set(
      ['send_signal', 'suspend', 'resume', 'terminate', 'kill', 'wait',
       'is_running', 'as_dict', 'parent', 'parents', 'children', 'rlimit',
       'memory_info_ex', 'oneshot']])
+if MACOS:
+    _as_dict_attrnames.remove('ionice')
 
 
 # =====================================================================
diff --git a/psutil/_psosx.py b/psutil/_psosx.py
index 7459a0f3..a25ee667 100644
--- a/psutil/_psosx.py
+++ b/psutil/_psosx.py
@@ -8,6 +8,7 @@ import contextlib
 import errno
 import functools
 import os
+import sys
 from socket import AF_INET
 from collections import namedtuple
 
@@ -24,9 +25,16 @@ from ._common import sockfam_to_enum
 from ._common import socktype_to_enum
 from ._common import usage_percent
 
+if sys.version_info >= (3, 4):
+    import enum
+else:
+    enum = None
 
-__extra__all__ = []
-
+__extra__all__ = [
+    # IO priority constants
+    "IOPOL_DEFAULT", "IOPOL_IMPORTANT", "IOPOL_PASSIVE", "IOPOL_THROTTLE",
+    "IOPOL_UTILITY", "IOPOL_STANDARD"
+]
 
 # =====================================================================
 # --- globals
@@ -59,6 +67,26 @@ PROC_STATUSES = {
     cext.SZOMB: _common.STATUS_ZOMBIE,
 }
 
+# [get|set]iopolicy constants (ionice() / I/O priority)
+# http://www.manpagez.com/man/3/setiopolicy_np/
+if enum is None:
+    IOPOL_DEFAULT = cext.IOPOL_DEFAULT
+    IOPOL_IMPORTANT = cext.IOPOL_IMPORTANT
+    IOPOL_PASSIVE = cext.IOPOL_PASSIVE
+    IOPOL_THROTTLE = cext.IOPOL_THROTTLE
+    IOPOL_UTILITY = cext.IOPOL_UTILITY
+    IOPOL_STANDARD = cext.IOPOL_STANDARD
+else:
+    class IOPriority(enum.IntEnum):
+        IOPOL_DEFAULT = cext.IOPOL_DEFAULT
+        IOPOL_IMPORTANT = cext.IOPOL_IMPORTANT
+        IOPOL_PASSIVE = cext.IOPOL_PASSIVE
+        IOPOL_THROTTLE = cext.IOPOL_THROTTLE
+        IOPOL_UTILITY = cext.IOPOL_UTILITY
+        IOPOL_STANDARD = cext.IOPOL_STANDARD
+
+    globals().update(IOPriority.__members__)
+
 kinfo_proc_map = dict(
     ppid=0,
     ruid=1,
@@ -576,3 +604,21 @@ class Process(object):
             ntuple = _common.pthread(thread_id, utime, stime)
             retlist.append(ntuple)
         return retlist
+
+    @wrap_exceptions
+    def ionice_get(self):
+        if self.pid != os.getpid():
+            raise ValueError("can't set I/O priority for another process")
+        policy = cext.proc_ioprio_get(self.pid)
+        if enum is not None:
+            policy = IOPriority(policy)
+        return policy
+
+    @wrap_exceptions
+    def ionice_set(self, policy, _):
+        if self.pid != os.getpid():
+            raise ValueError("can't set I/O priority for another process")
+        if _:
+            raise TypeError("set_proc_ionice() on macOS takes only "
+                            "1 argument (2 given)")
+        return cext.proc_ioprio_set(policy)
diff --git a/psutil/_psutil_osx.c b/psutil/_psutil_osx.c
index eaf4e514..5e898118 100644
--- a/psutil/_psutil_osx.c
+++ b/psutil/_psutil_osx.c
@@ -1743,6 +1743,37 @@ error:
 
 
 /*
+ * Get current process I/O priority (not possible for other PIDs).
+ */
+static PyObject *
+psutil_proc_ioprio_get(PyObject *self, PyObject *args) {
+    int ret;
+
+    ret = getiopolicy_np(IOPOL_TYPE_DISK, IOPOL_SCOPE_PROCESS);
+    if (ret == -1)
+        return PyErr_SetFromErrno(PyExc_OSError);
+    return Py_BuildValue("i", ret);
+}
+
+
+/*
+ * Set current process I/O priority (not possible for other PIDs).
+ */
+static PyObject *
+psutil_proc_ioprio_set(PyObject *self, PyObject *args) {
+    int ret;
+    int policy;
+
+    if (! PyArg_ParseTuple(args, "i", &policy))
+        return NULL;
+    ret = setiopolicy_np(IOPOL_TYPE_DISK, IOPOL_SCOPE_PROCESS, policy);
+    if (ret == -1)
+        return PyErr_SetFromErrno(PyExc_OSError);
+    Py_RETURN_NONE;
+}
+
+
+/*
  * define the psutil C module methods and initialize the module.
  */
 static PyMethodDef
@@ -1773,6 +1804,10 @@ PsutilMethods[] = {
      "Return the number of fds opened by process."},
     {"proc_connections", psutil_proc_connections, METH_VARARGS,
      "Get process TCP and UDP connections as a list of tuples"},
+    {"proc_ioprio_get", psutil_proc_ioprio_get, METH_VARARGS,
+     "Get current process' I/O priority."},
+    {"proc_ioprio_set", psutil_proc_ioprio_set, METH_VARARGS,
+     "Set current process' I/O priority."},
 
     // --- system-related functions
 
@@ -1897,6 +1932,13 @@ init_psutil_osx(void)
     PyModule_AddIntConstant(module, "TCPS_TIME_WAIT", TCPS_TIME_WAIT);
     PyModule_AddIntConstant(module, "PSUTIL_CONN_NONE", PSUTIL_CONN_NONE);
 
+    PyModule_AddIntConstant(module, "IOPOL_DEFAULT", IOPOL_DEFAULT);
+    PyModule_AddIntConstant(module, "IOPOL_IMPORTANT", IOPOL_IMPORTANT);
+    PyModule_AddIntConstant(module, "IOPOL_PASSIVE", IOPOL_PASSIVE);
+    PyModule_AddIntConstant(module, "IOPOL_THROTTLE", IOPOL_THROTTLE);
+    PyModule_AddIntConstant(module, "IOPOL_UTILITY", IOPOL_UTILITY);
+    PyModule_AddIntConstant(module, "IOPOL_STANDARD", IOPOL_STANDARD);
+
     // Exception.
     ZombieProcessError = PyErr_NewException(
         "_psutil_osx.ZombieProcessError", NULL, NULL);
diff --git a/psutil/tests/test_contracts.py b/psutil/tests/test_contracts.py
index 08e9e9b8..d39213cb 100755
--- a/psutil/tests/test_contracts.py
+++ b/psutil/tests/test_contracts.py
@@ -12,6 +12,7 @@ Some of these are duplicates of tests test_system.py and test_process.py
 import errno
 import os
 import stat
+import sys
 import time
 import traceback
 import warnings
@@ -47,6 +48,10 @@ from psutil.tests import VALID_PROC_STATUSES
 from psutil.tests import warn
 import psutil
 
+if sys.version_info >= (3, 4):
+    import enum
+else:
+    enum = None
 
 # ===================================================================
 # --- APIs availability
@@ -142,7 +147,8 @@ class TestAvailability(unittest.TestCase):
         self.assertEqual(hasattr(psutil.Process, "terminal"), POSIX)
 
     def test_proc_ionice(self):
-        self.assertEqual(hasattr(psutil.Process, "ionice"), LINUX or WINDOWS)
+        self.assertEqual(hasattr(psutil.Process, "ionice"),
+                         LINUX or WINDOWS or MACOS)
 
     def test_proc_rlimit(self):
         self.assertEqual(hasattr(psutil.Process, "rlimit"), LINUX)
@@ -334,6 +340,8 @@ class TestFetchAllProcesses(unittest.TestCase):
         default = object()
         failures = []
         for p, name in self.iter_procs():
+            if name == 'ionice' and MACOS and p.pid != os.getpid():
+                continue
             ret = default
             try:
                 ret = self.call_meth(p, name)
@@ -457,16 +465,19 @@ class TestFetchAllProcesses(unittest.TestCase):
                 self.assertGreaterEqual(field, 0)
 
     def ionice(self, ret, proc):
-        if POSIX:
+        if LINUX:
             assert is_namedtuple(ret)
             for field in ret:
                 self.assertIsInstance(field, int)
-        if LINUX:
             self.assertGreaterEqual(ret.ioclass, 0)
             self.assertGreaterEqual(ret.value, 0)
         else:
             self.assertGreaterEqual(ret, 0)
-            self.assertIn(ret, (0, 1, 2))
+            if enum is not None:
+                self.assertIn(
+                    ret, psutil._psplatform.IOPriority.__members__.values())
+            else:
+                self.assertIn(ret, (0, 1, 2))
 
     def num_threads(self, ret, proc):
         self.assertIsInstance(ret, int)
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index a550809f..f38803a0 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -353,6 +353,7 @@ class TestProcess(unittest.TestCase):
     @unittest.skipIf(not HAS_IONICE, "not supported")
     @unittest.skipIf(WINDOWS and get_winver() < WIN_VISTA, 'not supported')
     def test_ionice(self):
+        p = psutil.Process()
         if LINUX:
             from psutil import (IOPRIO_CLASS_NONE, IOPRIO_CLASS_RT,
                                 IOPRIO_CLASS_BE, IOPRIO_CLASS_IDLE)
@@ -360,7 +361,6 @@ class TestProcess(unittest.TestCase):
             self.assertEqual(IOPRIO_CLASS_RT, 1)
             self.assertEqual(IOPRIO_CLASS_BE, 2)
             self.assertEqual(IOPRIO_CLASS_IDLE, 3)
-            p = psutil.Process()
             try:
                 p.ionice(2)
                 ioclass, value = p.ionice()
@@ -384,8 +384,21 @@ class TestProcess(unittest.TestCase):
                 self.assertEqual(value, 7)
             finally:
                 p.ionice(IOPRIO_CLASS_NONE)
+        if MACOS:
+            names = ("IOPOL_IMPORTANT", "IOPOL_STANDARD", "IOPOL_UTILITY",
+                     "IOPOL_THROTTLE", "IOPOL_PASSIVE")
+            original = p.ionice()
+            try:
+                for name in names:
+                    if name == "IOPOL_IMPORTANT":
+                        # XXX has no effect (?!?)
+                        continue
+                    value = getattr(psutil, name)
+                    p.ionice(value)
+                    self.assertEqual(p.ionice(), value)
+            finally:
+                p.ionice(original)
         else:
-            p = psutil.Process()
             original = p.ionice()
             self.assertIsInstance(original, int)
             try:
@@ -417,8 +430,10 @@ class TestProcess(unittest.TestCase):
                 ValueError, "'ioclass' argument must be specified",
                 p.ionice, value=1)
         else:
-            self.assertRaises(ValueError, p.ionice, 3)
-            self.assertRaises(TypeError, p.ionice, 2, 1)
+            if not MACOS:
+                # ionice() can only be set for the current process
+                self.assertRaises(TypeError, p.ionice, 2, 1)
+                self.assertRaises(ValueError, p.ionice, 3)
 
     @unittest.skipIf(not HAS_RLIMIT, "not supported")
     def test_rlimit_get(self):
@@ -1280,6 +1295,9 @@ class TestProcess(unittest.TestCase):
                     else:
                         ret = meth(psutil.NORMAL_PRIORITY_CLASS)
                 elif name == 'ionice':
+                    if MACOS:
+                        # ionice() on OSX can only be used for os.getpid()
+                        continue
                     ret = meth()
                     ret = meth(2)
                 elif name == 'rlimit':
@@ -1349,7 +1367,7 @@ class TestProcess(unittest.TestCase):
                     raise
 
         succeed_or_zombie_p_exc(zproc.nice, 0)
-        if hasattr(zproc, 'ionice'):
+        if hasattr(zproc, 'ionice') and not MACOS:
             if LINUX:
                 succeed_or_zombie_p_exc(zproc.ionice, 2, 0)
             else: