Address code review

corona10 · Oct 1, 2023 · 2f0dc1c · 2f0dc1c
1 parent 936c182
commit 2f0dc1c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/Lib/test/test_cmd_line.py b/Lib/test/test_cmd_line.py
@@ -912,6 +912,8 @@ def test_cpu_count_default(self):
         self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
         res = assert_python_ok('-X', 'cpu_count=default', '-c', code, PYTHONCPUCOUNT='1234')
         self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
+        es = assert_python_ok('-c', code, PYTHONCPUCOUNT='default')
+        self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
 
     def res2int(self, res):
         out = res.out.strip().decode("utf-8")

diff --git a/Python/initconfig.c b/Python/initconfig.c
@@ -231,7 +231,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: override CPU count of os.cpu_count().\n\
+-X cpu_count=[n|default]: override CPU count of os.cpu_count() and os.process_cpu_count().\n\
     This helps for users who need to limit CPU resources of a container system."
 
 #ifdef Py_STATS
@@ -1628,9 +1628,12 @@ config_init_cpu_count(PyConfig *config)
     const char *env = config_get_env(config, "PYTHONCPUCOUNT");
     if (env) {
         int cpu_count = -1;
-        if (_Py_str_to_int(env, &cpu_count) != 0) {
+        if (strcmp(env, "default") == 0) {
             cpu_count = -1;
         }
+        else if (_Py_str_to_int(env, &cpu_count) != 0) {
+            goto error;
+        }
         if (cpu_count >= 1) {
             config->cpu_count = cpu_count;
         }