Issue #17914: Add os.cpu_count(). Patch by Yogesh Chaudhari, based on an

initial patch by Trent Nelson.

6 files changed, 94 insertions, 26 deletions

diff --git a/Doc/library/multiprocessing.rst b/Doc/library/multiprocessing.rst
index 1c18062..2562e30 100644
--- a/Doc/library/multiprocessing.rst
+++ b/Doc/library/multiprocessing.rst
@@ -702,6 +702,9 @@ Miscellaneous
    Return the number of CPUs in the system.  May raise
    :exc:`NotImplementedError`.
 
+   .. seealso::
+      :func:`os.cpu_count`
+
 .. function:: current_process()
 
    Return the :class:`Process` object corresponding to the current process.
diff --git a/Doc/library/os.rst b/Doc/library/os.rst
index 4ab5ebd..137fd42 100644
--- a/Doc/library/os.rst
+++ b/Doc/library/os.rst
@@ -3155,10 +3155,6 @@ operating system.
    Return the set of CPUs the process with PID *pid* (or the current process
    if zero) is restricted to.
 
-   .. seealso::
-      :func:`multiprocessing.cpu_count` returns the number of CPUs in the
-      system.
-
 
 .. _os-path:
 
@@ -3196,6 +3192,13 @@ Miscellaneous System Information
    Availability: Unix.
 
 
+.. function:: cpu_count()
+
+   Return the number of CPUs in the system. Returns None if undetermined.
+
+   .. versionadded:: 3.4
+
+
 .. function:: getloadavg()
 
    Return the number of processes in the system run queue averaged over the last
diff --git a/Lib/multiprocessing/__init__.py b/Lib/multiprocessing/__init__.py
index b5f16d7..b42613f 100644
--- a/Lib/multiprocessing/__init__.py
+++ b/Lib/multiprocessing/__init__.py
@@ -85,30 +85,11 @@ def cpu_count():
     '''
     Returns the number of CPUs in the system
     '''
-    if sys.platform == 'win32':
-        try:
-            num = int(os.environ['NUMBER_OF_PROCESSORS'])
-        except (ValueError, KeyError):
-            num = 0
-    elif 'bsd' in sys.platform or sys.platform == 'darwin':
-        comm = '/sbin/sysctl -n hw.ncpu'
-        if sys.platform == 'darwin':
-            comm = '/usr' + comm
-        try:
-            with os.popen(comm) as p:
-                num = int(p.read())
-        except ValueError:
-            num = 0
+    num = os.cpu_count()
+    if num is None:
+        raise NotImplementedError('cannot determine number of cpus')
     else:
-        try:
-            num = os.sysconf('SC_NPROCESSORS_ONLN')
-        except (ValueError, OSError, AttributeError):
-            num = 0
-
-    if num >= 1:
         return num
-    else:
-        raise NotImplementedError('cannot determine number of cpus')
 
 def freeze_support():
     '''
diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py
index 3a38285..65d3c3b 100644
--- a/Lib/test/test_os.py
+++ b/Lib/test/test_os.py
@@ -2216,6 +2216,15 @@ class OSErrorTests(unittest.TestCase):
                 else:
                     self.fail("No exception thrown by {}".format(func))
 
+class CPUCountTests(unittest.TestCase):
+    def test_cpu_count(self):
+        cpus = os.cpu_count()
+        if cpus is not None:
+            self.assertIsInstance(cpus, int)
+            self.assertGreater(cpus, 0)
+        else:
+            self.skipTest("Could not determine the number of CPUs")
+
 @support.reap_threads
 def test_main():
     support.run_unittest(
@@ -2246,6 +2255,7 @@ def test_main():
         TermsizeTests,
         OSErrorTests,
         RemoveDirsTests,
+        CPUCountTests,
     )
 
 if __name__ == "__main__":
diff --git a/Misc/NEWS b/Misc/NEWS
index 1c54d33..541a424 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -99,6 +99,9 @@ Core and Builtins
 Library
 -------
 
+- Issue #17914: Add os.cpu_count(). Patch by Yogesh Chaudhari, based on an
+  initial patch by Trent Nelson.
+
 - Issue #17980: Fix possible abuse of ssl.match_hostname() for denial of
   service using certificates with many wildcards (CVE-2013-2099).
 
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
index cb72bf0..c24e162 100644
--- a/Modules/posixmodule.c
+++ b/Modules/posixmodule.c
@@ -113,6 +113,18 @@ corresponding Unix manual entries for more information on calls.");
 #include <dlfcn.h>
 #endif
 
+#ifdef __hpux
+#include <sys/mpctl.h>
+#endif
+
+#if defined(__DragonFly__) || \
+    defined(__OpenBSD__)   || \
+    defined(__FreeBSD__)   || \
+    defined(__NetBSD__)    || \
+    defined(__APPLE__)
+#include <sys/sysctl.h>
+#endif
+
 #if defined(MS_WINDOWS)
 #  define TERMSIZE_USE_CONIO
 #elif defined(HAVE_SYS_IOCTL_H)
@@ -10302,6 +10314,60 @@ get_terminal_size(PyObject *self, PyObject *args)
 }
 #endif /* defined(TERMSIZE_USE_CONIO) || defined(TERMSIZE_USE_IOCTL) */
 
+PyDoc_STRVAR(posix_cpu_count__doc__,
+"cpu_count() -> integer\n\n\
+Return the number of CPUs in the system, or None if this value cannot be\n\
+established.");
+
+#if defined(__DragonFly__) || \
+    defined(__OpenBSD__)   || \
+    defined(__FreeBSD__)   || \
+    defined(__NetBSD__)    || \
+    defined(__APPLE__)
+static long
+_bsd_cpu_count(void)
+{
+    long ncpu = 0;
+    int mib[2];
+    size_t len = sizeof(int);
+
+    mib[0] = CTL_HW;
+    mib[1] = HW_NCPU;
+    if (sysctl(mib, 2, &ncpu, &len, NULL, 0) == 0)
+        return ncpu;
+    else
+        return 0;
+}
+#endif
+
+static PyObject *
+posix_cpu_count(PyObject *self)
+{
+    long ncpu = 0;
+#ifdef MS_WINDOWS
+    SYSTEM_INFO sysinfo;
+    GetSystemInfo(&sysinfo);
+    ncpu = sysinfo.dwNumberOfProcessors;
+#elif defined(__hpux)
+    ncpu = mpctl(MPC_GETNUMSPUS, NULL, NULL);
+#elif defined(HAVE_SYSCONF) && defined(_SC_NPROCESSORS_ONLN)
+    ncpu = sysconf(_SC_NPROCESSORS_ONLN);
+#elif defined(__APPLE__)
+    size_t len = sizeof(int);
+    if (sysctlnametomib("hw.logicalcpu", &ncpu, &len, NULL, 0) != 0)
+        ncpu = _bsd_cpu_count();
+#elif defined(__DragonFly__) || \
+      defined(__OpenBSD__)   || \
+      defined(__FreeBSD__)   || \
+      defined(__NetBSD__)
+    ncpu = _bsd_cpu_count();
+#endif
+    if (ncpu >= 1)
+        return PyLong_FromLong(ncpu);
+    else
+        Py_RETURN_NONE;
+}
+
 
 static PyMethodDef posix_methods[] = {
     {"access",          (PyCFunction)posix_access,
@@ -10747,6 +10813,8 @@ static PyMethodDef posix_methods[] = {
 #if defined(TERMSIZE_USE_CONIO) || defined(TERMSIZE_USE_IOCTL)
     {"get_terminal_size", get_terminal_size, METH_VARARGS, termsize__doc__},
 #endif
+    {"cpu_count", (PyCFunction)posix_cpu_count,
+                  METH_NOARGS, posix_cpu_count__doc__},
     {NULL,              NULL}            /* Sentinel */
 };