From 070d8665670e929ddf26a217d2c1a4116f82d30e Mon Sep 17 00:00:00 2001
From: Victor Stinner <vstinner@python.org>
Date: Fri, 25 Apr 2025 18:29:55 +0200
Subject: gh-107954: Allow setting cpu_count in PyConfig_Set() (#132954)

* gh-107954: Allow setting cpu_count in PyConfig_Set()

* Update the doc
---
 Doc/c-api/init_config.rst         |  2 +-
 Lib/test/test_capi/test_config.py | 14 +++++++++++++-
 Python/initconfig.c               | 30 +++++++++++++++++++++++-------
 3 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/Doc/c-api/init_config.rst b/Doc/c-api/init_config.rst
index 74f34f8..d8d442c 100644
--- a/Doc/c-api/init_config.rst
+++ b/Doc/c-api/init_config.rst
@@ -320,7 +320,7 @@ Configuration Options
    * - ``"cpu_count"``
      - :c:member:`cpu_count <PyConfig.cpu_count>`
      - ``int``
-     - Read-only
+     - Public
    * - ``"dev_mode"``
      - :c:member:`dev_mode <PyConfig.dev_mode>`
      - ``bool``
diff --git a/Lib/test/test_capi/test_config.py b/Lib/test/test_capi/test_config.py
index 6d27ec5..bf351c4 100644
--- a/Lib/test/test_capi/test_config.py
+++ b/Lib/test/test_capi/test_config.py
@@ -1,6 +1,7 @@
 """
 Tests PyConfig_Get() and PyConfig_Set() C API (PEP 741).
 """
+import os
 import sys
 import sysconfig
 import types
@@ -374,12 +375,23 @@ class CAPITests(unittest.TestCase):
                 finally:
                     config_set(name, old_value)
 
+    def test_config_set_cpu_count(self):
+        config_get = _testcapi.config_get
+        config_set = _testcapi.config_set
+
+        old_value = config_get('cpu_count')
+        try:
+            config_set('cpu_count', 123)
+            self.assertEqual(os.cpu_count(), 123)
+        finally:
+            config_set('cpu_count', old_value)
+
     def test_config_set_read_only(self):
         # Test PyConfig_Set() on read-only options
         config_set = _testcapi.config_set
         for name, value in (
             ("allocator", 0),  # PyPreConfig member
-            ("cpu_count", 8),
+            ("perf_profiling", 8),
             ("dev_mode", True),
             ("filesystem_encoding", "utf-8"),
         ):
diff --git a/Python/initconfig.c b/Python/initconfig.c
index 39e70bd..43b85a3 100644
--- a/Python/initconfig.c
+++ b/Python/initconfig.c
@@ -110,6 +110,7 @@ static const PyConfigSpec PYCONFIG_SPEC[] = {
     SPEC(base_executable, WSTR_OPT, PUBLIC, SYS_ATTR("_base_executable")),
     SPEC(base_prefix, WSTR_OPT, PUBLIC, SYS_ATTR("base_prefix")),
     SPEC(bytes_warning, UINT, PUBLIC, SYS_FLAG(9)),
+    SPEC(cpu_count, INT, PUBLIC, NO_SYS),
     SPEC(exec_prefix, WSTR_OPT, PUBLIC, SYS_ATTR("exec_prefix")),
     SPEC(executable, WSTR_OPT, PUBLIC, SYS_ATTR("executable")),
     SPEC(inspect, BOOL, PUBLIC, SYS_FLAG(1)),
@@ -138,7 +139,6 @@ static const PyConfigSpec PYCONFIG_SPEC[] = {
     SPEC(check_hash_pycs_mode, WSTR, READ_ONLY, NO_SYS),
     SPEC(code_debug_ranges, BOOL, READ_ONLY, NO_SYS),
     SPEC(configure_c_stdio, BOOL, READ_ONLY, NO_SYS),
-    SPEC(cpu_count, INT, READ_ONLY, NO_SYS),
     SPEC(dev_mode, BOOL, READ_ONLY, NO_SYS),  // sys.flags.dev_mode
     SPEC(dump_refs, BOOL, READ_ONLY, NO_SYS),
     SPEC(dump_refs_file, WSTR_OPT, READ_ONLY, NO_SYS),
@@ -4524,6 +4524,18 @@ error:
 }
 
 
+// Set PyConfig.ATTR integer member
+static int
+config_set_int_attr(const PyConfigSpec *spec, int value)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    PyConfig *config = &interp->config;
+    int *member = config_get_spec_member(config, spec);
+    *member = value;
+    return 0;
+}
+
+
 int
 PyConfig_Set(const char *name, PyObject *value)
 {
@@ -4632,17 +4644,21 @@ PyConfig_Set(const char *name, PyObject *value)
         Py_UNREACHABLE();
     }
 
-
     if (spec->sys.attr != NULL) {
         // Set the sys attribute, but don't set PyInterpreterState.config
         // to keep the code simple.
         return PySys_SetObject(spec->sys.attr, value);
     }
-    else if (spec->sys.flag_index >= 0 && has_int_value) {
-        return config_set_sys_flag(spec, int_value);
-    }
-    else if (strcmp(spec->name, "int_max_str_digits") == 0 && has_int_value) {
-        return _PySys_SetIntMaxStrDigits(int_value);
+    else if (has_int_value) {
+        if (spec->sys.flag_index >= 0) {
+            return config_set_sys_flag(spec, int_value);
+        }
+        else if (strcmp(spec->name, "int_max_str_digits") == 0) {
+            return _PySys_SetIntMaxStrDigits(int_value);
+        }
+        else {
+            return config_set_int_attr(spec, int_value);
+        }
     }
 
 cannot_set:
-- 
cgit v0.12