OSX: implement sensors_battery (#1177)

6 files changed, 129 insertions, 3 deletions

diff --git a/psutil/__init__.py b/psutil/__init__.py
index 12cd2c4c..c8da0a3d 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -2210,7 +2210,7 @@ if hasattr(_psplatform, "sensors_fans"):
     __all__.append("sensors_fans")
 
 
-# Linux, Windows, FreeBSD
+# Linux, Windows, FreeBSD, OSX
 if hasattr(_psplatform, "sensors_battery"):
 
     def sensors_battery():
diff --git a/psutil/_psosx.py b/psutil/_psosx.py
index e9b6ed82..9fa7716d 100644
--- a/psutil/_psosx.py
+++ b/psutil/_psosx.py
@@ -208,6 +208,29 @@ def disk_partitions(all=False):
 
 
 # =====================================================================
+# --- sensors
+# =====================================================================
+
+
+def sensors_battery():
+    """Return battery information.
+    """
+    try:
+        percent, minsleft, power_plugged = cext.sensors_battery()
+    except NotImplementedError:
+        # no power source - return None according to interface
+        return None
+    power_plugged = power_plugged == 1
+    if power_plugged:
+        secsleft = _common.POWER_TIME_UNLIMITED
+    elif minsleft == -1:
+        secsleft = _common.POWER_TIME_UNKNOWN
+    else:
+        secsleft = minsleft * 60
+    return _common.sbattery(percent, secsleft, power_plugged)
+
+
+# =====================================================================
 # --- network
 # =====================================================================
 
diff --git a/psutil/_psutil_osx.c b/psutil/_psutil_osx.c
index 46751626..6c520e5d 100644
--- a/psutil/_psutil_osx.c
+++ b/psutil/_psutil_osx.c
@@ -38,6 +38,8 @@
 #include <IOKit/storage/IOBlockStorageDriver.h>
 #include <IOKit/storage/IOMedia.h>
 #include <IOKit/IOBSD.h>
+#include <IOKit/ps/IOPowerSources.h>
+#include <IOKit/ps/IOPSKeys.h>
 
 #include "_psutil_common.h"
 #include "_psutil_posix.h"
@@ -1788,6 +1790,92 @@ psutil_cpu_stats(PyObject *self, PyObject *args) {
 }
 
 
+/*
+ * Return battery information.
+ */
+static PyObject *
+psutil_sensors_battery(PyObject *self, PyObject *args) {
+    PyObject *py_tuple = NULL;
+    CFTypeRef power_info = NULL;
+    CFArrayRef power_sources_list = NULL;
+    CFDictionaryRef power_sources_information = NULL;
+    CFNumberRef capacity_ref = NULL;
+    CFNumberRef time_to_empty_ref = NULL;
+    CFStringRef ps_state_ref = NULL;
+    uint32_t capacity;     /* units are percent */
+    int time_to_empty;     /* units are minutes */
+    int is_power_plugged;
+
+    power_info = IOPSCopyPowerSourcesInfo();
+
+    if (!power_info) {
+        PyErr_SetString(PyExc_RuntimeError,
+            "IOPSCopyPowerSourcesInfo() syscall failed");
+        goto error;
+    }
+
+    power_sources_list = IOPSCopyPowerSourcesList(power_info);
+    if (!power_sources_list) {
+        PyErr_SetString(PyExc_RuntimeError,
+            "IOPSCopyPowerSourcesList() syscall failed");
+        goto error;
+    }
+
+    /* Should only get one source. But in practice, check for > 0 sources */
+    if (!CFArrayGetCount(power_sources_list)) {
+        PyErr_SetString(PyExc_NotImplementedError, "no battery");
+        goto error;
+    }
+
+    power_sources_information = IOPSGetPowerSourceDescription(
+        power_info, CFArrayGetValueAtIndex(power_sources_list, 0));
+
+    capacity_ref = (CFNumberRef)  CFDictionaryGetValue(
+        power_sources_information, CFSTR(kIOPSCurrentCapacityKey));
+    if (!CFNumberGetValue(capacity_ref, kCFNumberSInt32Type, &capacity)) {
+        PyErr_SetString(PyExc_RuntimeError,
+            "No battery capacity infomration in power sources info");
+        goto error;
+    }
+
+    ps_state_ref = (CFStringRef) CFDictionaryGetValue(
+        power_sources_information, CFSTR(kIOPSPowerSourceStateKey));
+    is_power_plugged = CFStringCompare(
+        ps_state_ref, CFSTR(kIOPSACPowerValue), 0)
+        == kCFCompareEqualTo;
+
+    time_to_empty_ref = (CFNumberRef) CFDictionaryGetValue(
+        power_sources_information, CFSTR(kIOPSTimeToEmptyKey));
+    if (!CFNumberGetValue(time_to_empty_ref,
+                          kCFNumberIntType, &time_to_empty)) {
+        /* This value is recommended for non-Apple power sources, so it's not
+         * an error if it doesn't exist. We'll return -1 for "unknown" */
+        /* A value of -1 indicates "Still Calculating the Time" also for
+         * apple power source */
+        time_to_empty = -1;
+    }
+
+    py_tuple = Py_BuildValue("Iii",
+        capacity, time_to_empty, is_power_plugged);
+    if (!py_tuple) {
+        goto error;
+    }
+
+    CFRelease(power_info);
+    CFRelease(power_sources_list);
+    /* Caller should NOT release power_sources_information */
+
+    return py_tuple;
+
+error:
+    if (power_info)
+        CFRelease(power_info);
+    if (power_sources_list)
+        CFRelease(power_sources_list);
+    Py_XDECREF(py_tuple);
+    return NULL;
+}
+
 
 /*
  * define the psutil C module methods and initialize the module.
@@ -1854,6 +1942,8 @@ PsutilMethods[] = {
      "Return currently connected users as a list of tuples"},
     {"cpu_stats", psutil_cpu_stats, METH_VARARGS,
      "Return CPU statistics"},
+    {"sensors_battery", psutil_sensors_battery, METH_VARARGS,
+     "Return battery information."},
 
     // --- others
     {"set_testing", psutil_set_testing, METH_NOARGS,
diff --git a/psutil/tests/__init__.py b/psutil/tests/__init__.py
index 4f7995dc..7c58ab3c 100644
--- a/psutil/tests/__init__.py
+++ b/psutil/tests/__init__.py
@@ -70,7 +70,7 @@ __all__ = [
     'VERBOSITY',
     "HAS_CPU_AFFINITY", "HAS_CPU_FREQ", "HAS_ENVIRON", "HAS_PROC_IO_COUNTERS",
     "HAS_IONICE", "HAS_MEMORY_MAPS", "HAS_PROC_CPU_NUM", "HAS_RLIMIT",
-    "HAS_SENSORS_BATTERY", "HAS_BATTERY""HAS_SENSORS_FANS",
+    "HAS_SENSORS_BATTERY", "HAS_BATTERY", "HAS_SENSORS_FANS",
     "HAS_SENSORS_TEMPERATURES", "HAS_MEMORY_FULL_INFO",
     # subprocesses
     'pyrun', 'reap_children', 'get_test_subprocess', 'create_zombie_proc',
diff --git a/psutil/tests/test_contracts.py b/psutil/tests/test_contracts.py
index 5e5c2e9a..4c57b25a 100755
--- a/psutil/tests/test_contracts.py
+++ b/psutil/tests/test_contracts.py
@@ -123,7 +123,7 @@ class TestAvailability(unittest.TestCase):
 
     def test_battery(self):
         self.assertEqual(hasattr(psutil, "sensors_battery"),
-                         LINUX or WINDOWS or FREEBSD)
+                         LINUX or WINDOWS or FREEBSD or OSX)
 
     def test_proc_environ(self):
         self.assertEqual(hasattr(psutil.Process, "environ"),
diff --git a/psutil/tests/test_osx.py b/psutil/tests/test_osx.py
index c8214f14..bcb2ba4e 100755
--- a/psutil/tests/test_osx.py
+++ b/psutil/tests/test_osx.py
@@ -14,6 +14,7 @@ import psutil
 from psutil import OSX
 from psutil.tests import create_zombie_proc
 from psutil.tests import get_test_subprocess
+from psutil.tests import HAS_BATTERY
 from psutil.tests import MEMORY_TOLERANCE
 from psutil.tests import reap_children
 from psutil.tests import retry_before_failing
@@ -285,6 +286,18 @@ class TestSystemAPIs(unittest.TestCase):
                 self.assertEqual(stats.mtu,
                                  int(re.findall(r'mtu (\d+)', out)[0]))
 
+    # --- sensors_battery
+
+    @unittest.skipIf(not HAS_BATTERY, "no battery")
+    def test_sensors_battery(self):
+        out = sh("pmset -g batt")
+        percent = re.search("(\d+)%", out).group(1)
+        drawing_from = re.search("Now drawing from '([^']+)'", out).group(1)
+        power_plugged = drawing_from == "AC Power"
+        psutil_result = psutil.sensors_battery()
+        self.assertEqual(psutil_result.power_plugged, power_plugged)
+        self.assertEqual(psutil_result.percent, int(percent))
+
 
 if __name__ == '__main__':
     run_test_module_by_name(__file__)