gh-109595: Update -Xcpu_count=<n> cmdline to support process mode

corona10 · corona10 · commit 920fb835f19f · 2023-10-10T23:37:44.000+09:00
diff --git a/Doc/c-api/init_config.rst b/Doc/c-api/init_config.rst
@@ -884,7 +884,7 @@ PyConfig
       override the return values of :func:`os.cpu_count`,
       :func:`os.process_cpu_count`, and :func:`multiprocessing.cpu_count`.
 
-      Configured by the :samp:`-X cpu_count={n|default}` command line
+      Configured by the :samp:`-X cpu_count={n|default|process}` command line
       flag or the :envvar:`PYTHON_CPU_COUNT` environment variable.
 
       Default: ``-1``.
diff --git a/Doc/using/cmdline.rst b/Doc/using/cmdline.rst
@@ -551,7 +551,9 @@ Miscellaneous options
      *n* must be greater than or equal to 1.
      This option may be useful for users who need to limit CPU resources of a
      container system. See also :envvar:`PYTHON_CPU_COUNT`.
-     If *n* is ``default``, nothing is overridden.
+     If *n* is ``default``, nothing is overridden, if *n* is ``process``,
+     :func:`os.os.process_cpu_count` will be alias of :func:`os.cpu_count`
+     and :func:`multiprocessing.cpu_count`.
 
    It also allows passing arbitrary values and retrieving them through the
    :data:`sys._xoptions` dictionary.
@@ -1076,6 +1078,8 @@ conflict.
 
    If this variable is set to a positive integer, it overrides the return
    values of :func:`os.cpu_count` and :func:`os.process_cpu_count`.
+   if this variable is set to ``process``, :func:`os.process_cpu_count` becomes
+   alias of :func:`os.cpu_count` and :func:`multiprocessing.cpu_count`.
 
    See also the :option:`-X cpu_count <-X>` command-line option.
 
diff --git a/Lib/os.py b/Lib/os.py
@@ -1147,6 +1147,9 @@ def process_cpu_count():
         current process. Return None if indeterminable.
         """
         return len(sched_getaffinity(0))
+    if sys._get_cpu_count_config() == -2:
+        # Just an alias to process_count() (same docstring)
+        cpu_count = process_cpu_count
 else:
     # Just an alias to cpu_count() (same docstring)
     process_cpu_count = cpu_count
diff --git a/Lib/test/test_cmd_line.py b/Lib/test/test_cmd_line.py
@@ -915,6 +915,15 @@ def test_cpu_count_default(self):
         es = assert_python_ok('-c', code, PYTHON_CPU_COUNT='default')
         self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
 
+    def test_cpu_count_process(self):
+        code = "import os; print(os.cpu_count(), os.process_cpu_count())"
+        res = assert_python_ok('-X', 'cpu_count=process', '-c', code)
+        self.assertEqual(self.res2int(res), (os.process_cpu_count(), os.process_cpu_count()))
+        res = assert_python_ok('-X', 'cpu_count=default', '-c', code, PYTHON_CPU_COUNT='process')
+        self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
+        es = assert_python_ok('-c', code, PYTHON_CPU_COUNT='process')
+        self.assertEqual(self.res2int(res), (os.process_cpu_count(), os.process_cpu_count()))
+
     def res2int(self, res):
         out = res.out.strip().decode("utf-8")
         return tuple(int(i) for i in out.split())
diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-10-10-23-37-42.gh-issue-109595.jPp6s9.rst b/Misc/NEWS.d/next/Core and Builtins/2023-10-10-23-37-42.gh-issue-109595.jPp6s9.rst
@@ -0,0 +1,4 @@
+Update :option:`-X cpu_count <-X>` command line option to support
+``process`` mode that :func:`os.process_cpu_count` becomes an alias of
+:func:`os.cpu_count` and :func:`multiprocessing.cpu_count`. Patch by Donghee
+Na.
diff --git a/Python/initconfig.c b/Python/initconfig.c
@@ -232,7 +232,7 @@ The following implementation-specific options are available:\n\
     This helps avoid denial of service attacks when parsing untrusted data.\n\
     The default is sys.int_info.default_max_str_digits.  0 disables.\n\
 \n\
--X cpu_count=[n|default]: Override the return value of os.cpu_count(),\n\
+-X cpu_count=[n|default|process]: Override the return value of os.cpu_count(),\n\
     os.process_cpu_count(), and multiprocessing.cpu_count(). This can help users who need\n\
     to limit resources in a container."
 
@@ -1636,6 +1636,9 @@ config_init_cpu_count(PyConfig *config)
         if (strcmp(env, "default") == 0) {
             cpu_count = -1;
         }
+        else if(strcmp(env, "process") == 0) {
+            cpu_count = -2;
+        }
         else if (_Py_str_to_int(env, &cpu_count) < 0 || cpu_count < 1) {
             goto error;
         }
@@ -1650,6 +1653,9 @@ config_init_cpu_count(PyConfig *config)
             if (wcscmp(sep + 1, L"default") == 0) {
                 cpu_count = -1;
             }
+            else if (wcscmp(sep + 1, L"process") == 0) {
+                cpu_count = -2;
+            }
             else if (config_wstr_to_int(sep + 1, &cpu_count) < 0 || cpu_count < 1) {
                 goto error;
             }
