Move the function to _winapi

Co-Authored-by: Cody Maloney <cmaloney@users.noreply.github.com>

Author: vstinner

Lib/test/libregrtest/utils.py

@@ -14,9 +14,9 @@
 import textwrap
 from collections.abc import Callable
 try:
-    from _testcapi import get_process_memory_usage as _get_process_memory_usage
+    import _winapi
 except ImportError:
-    _get_process_memory_usage = None
+    _winapi = None
 
 from test import support
 from test.support import os_helper
@@ -758,13 +758,7 @@ def display_title(title):
     print(flush=True)
 
 
-def get_process_memory_usage(pid: int) -> int | None:
-    """
-    Get process memory usage in bytes.
-    """
-    if _get_process_memory_usage is not None:
-        return _get_process_memory_usage(pid)
-
+def _get_process_memory_usage_linux(pid: int) -> int | None:
     # Linux implementation: read the private memory in bytes from
     # /proc/pid/smaps.
     try:
@@ -784,3 +778,25 @@ def get_process_memory_usage(pid: int) -> int | None:
         return total
     except ProcessLookupError:
         return None
+
+
+def _get_process_memory_usage_windows(pid: int) -> int | None:
+    handle = _winapi.OpenProcess(_winapi.PROCESS_QUERY_LIMITED_INFORMATION,
+                                 False, pid)
+    try:
+       mem_info = _winapi.GetProcessMemoryInfo(handle)
+    finally:
+        _winapi.CloseHandle(handle)
+    return mem_info['WorkingSetSize']
+
+
+if _winapi is not None:
+    get_process_memory_usage = _get_process_memory_usage_windows
+elif sys.platform == 'linux':
+    get_process_memory_usage = _get_process_memory_usage_linux
+else:
+    def get_process_memory_usage(pid: int) -> int | None:
+        """
+        Get process memory usage in bytes.
+        """
+        return None

Modules/_testcapi/mem.c

@@ -1,9 +1,5 @@
 #include "parts.h"
 
-#ifdef MS_WINDOWS
-#  include <windows.h>
-#  include <psapi.h>              // GetProcessMemoryInfo()
-#endif
 #include <stddef.h>
 
 
@@ -688,34 +684,6 @@ tracemalloc_track_race(PyObject *self, PyObject *args)
 }
 
 
-#ifdef MS_WINDOWS
-// Get process memory usage in bytes.
-static PyObject *
-get_process_memory_usage(PyObject *self, PyObject *args)
-{
-    int pid;
-    if (!PyArg_ParseTuple(args, "i", &pid)) {
-        return NULL;
-    }
-
-    // Get WorkingSetSize from GetProcessMemoryInfo()
-    HANDLE handle = OpenProcess(PROCESS_ALL_ACCESS, FALSE, (DWORD)pid);
-    if (handle == NULL) {
-        return PyErr_SetFromWindowsErr(0);
-    }
-
-    PROCESS_MEMORY_COUNTERS pmc;
-    if (!GetProcessMemoryInfo(handle, &pmc, sizeof(pmc))) {
-        CloseHandle(handle);
-        return PyErr_SetFromWindowsErr(0);
-    }
-    CloseHandle(handle);
-
-    return PyLong_FromSize_t(pmc.WorkingSetSize);
-}
-#endif
-
-
 static PyMethodDef test_methods[] = {
     {"pymem_api_misuse",              pymem_api_misuse,              METH_NOARGS},
     {"pymem_buffer_overflow",         pymem_buffer_overflow,         METH_NOARGS},
@@ -730,9 +698,6 @@ static PyMethodDef test_methods[] = {
     {"test_pymem_setrawallocators",   test_pymem_setrawallocators,   METH_NOARGS},
     {"test_pyobject_new",             test_pyobject_new,             METH_NOARGS},
     {"test_pyobject_setallocators",   test_pyobject_setallocators,   METH_NOARGS},
-#ifdef MS_WINDOWS
-    {"get_process_memory_usage",      get_process_memory_usage,      METH_VARARGS},
-#endif
 
     // Tracemalloc tests
     {"tracemalloc_track",             tracemalloc_track,             METH_VARARGS},

Modules/_winapi.c

@@ -49,6 +49,13 @@
 #include <crtdbg.h>
 #include "winreparse.h"
 
+// PSAPI_VERSION=2 redirects GetProcessMemoryInfo() to
+// K32GetProcessMemoryInfo() in kernel32.dll, so we don't need to link
+// psapi.lib. See:
+// https://learn.microsoft.com/windows/win32/api/psapi/nf-psapi-getprocessmemoryinfo
+#define PSAPI_VERSION 2
+#include <psapi.h>                // GetProcessMemoryInfo()
+
 #if defined(MS_WIN32) && !defined(MS_WIN64)
 #define HANDLE_TO_PYNUM(handle) \
     PyLong_FromUnsignedLong((unsigned long) handle)
@@ -3080,6 +3087,57 @@ _winapi_ReportEvent_impl(PyObject *module, HANDLE handle,
 }
 
 
+
+/*[clinic input]
+_winapi.GetProcessMemoryInfo
+    handle: HANDLE
+    /
+
+Return the memory usage of the given process handle.
+[clinic start generated code]*/
+
+static PyObject *
+_winapi_GetProcessMemoryInfo_impl(PyObject *module, HANDLE handle)
+/*[clinic end generated code: output=00a5d09732e84120 input=cfa669907cfbcecb]*/
+{
+    PROCESS_MEMORY_COUNTERS pmc;
+    if (!GetProcessMemoryInfo(handle, &pmc, sizeof(pmc))) {
+        return PyErr_SetFromWindowsErr(0);
+    }
+
+    PyObject *result = PyDict_New();
+    if (result == NULL) {
+        return NULL;
+    }
+
+#define ADD(ATTR) \
+    do { \
+        PyObject *obj = PyLong_FromSize_t(pmc.ATTR); \
+        if (obj == NULL) { \
+            Py_DECREF(result); \
+            return NULL; \
+        } \
+        if (PyDict_SetItemString(result, #ATTR, obj) < 0) { \
+            Py_DECREF(obj); \
+            return NULL; \
+        } \
+        Py_DECREF(obj); \
+    } while (0)
+
+    ADD(PageFaultCount);
+    ADD(PeakWorkingSetSize);
+    ADD(WorkingSetSize);
+    ADD(QuotaPeakPagedPoolUsage);
+    ADD(QuotaPagedPoolUsage);
+    ADD(QuotaPeakNonPagedPoolUsage);
+    ADD(QuotaNonPagedPoolUsage);
+    ADD(PagefileUsage);
+    ADD(PeakPagefileUsage);
+
+    return result;
+}
+
+
 static PyMethodDef winapi_functions[] = {
     _WINAPI_CLOSEHANDLE_METHODDEF
     _WINAPI_CONNECTNAMEDPIPE_METHODDEF
@@ -3130,6 +3188,7 @@ static PyMethodDef winapi_functions[] = {
     _WINAPI__MIMETYPES_READ_WINDOWS_REGISTRY_METHODDEF
     _WINAPI_NEEDCURRENTDIRECTORYFOREXEPATH_METHODDEF
     _WINAPI_COPYFILE2_METHODDEF
+    _WINAPI_GETPROCESSMEMORYINFO_METHODDEF
     {NULL, NULL}
 };
 
@@ -3226,6 +3285,7 @@ static int winapi_exec(PyObject *m)
     WINAPI_CONSTANT(F_DWORD, PROCESS_ALL_ACCESS);
     WINAPI_CONSTANT(F_DWORD, SYNCHRONIZE);
     WINAPI_CONSTANT(F_DWORD, PROCESS_DUP_HANDLE);
+    WINAPI_CONSTANT(F_DWORD, PROCESS_QUERY_LIMITED_INFORMATION);
     WINAPI_CONSTANT(F_DWORD, SEC_COMMIT);
     WINAPI_CONSTANT(F_DWORD, SEC_IMAGE);
     WINAPI_CONSTANT(F_DWORD, SEC_LARGE_PAGES);