Content-Length: 6103 | pFad | http://github.com/python/cpython/pull/150267.diff

thub.com diff --git a/Lib/test/libregrtest/utils.py b/Lib/test/libregrtest/utils.py index 1b4cb96406d6f6..21b84f7555b771 100644 --- a/Lib/test/libregrtest/utils.py +++ b/Lib/test/libregrtest/utils.py @@ -12,7 +12,13 @@ import sysconfig import tempfile import textwrap +import types from collections.abc import Callable +_winapi: types.ModuleType | None +try: + import _winapi +except ImportError: + _winapi = None from test import support from test.support import os_helper @@ -754,10 +760,9 @@ def display_title(title): print(flush=True) -def get_process_memory_usage(pid: int) -> int | None: - """ - Read the private memory in bytes from /proc/pid/smaps. - """ +def _get_process_memory_usage_linux(pid: int) -> int | None: + # Linux implementation: read the private memory in bytes from + # /proc/pid/smaps. try: fp = open(f"/proc/{pid}/smaps", "rb") except OSError: @@ -775,3 +780,26 @@ 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: + assert _winapi is not None # to make mypy happy + 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 diff --git a/Modules/_winapi.c b/Modules/_winapi.c index ffa407b2f21f73..fc2c0890468a6b 100644 --- a/Modules/_winapi.c +++ b/Modules/_winapi.c @@ -49,6 +49,13 @@ #include #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 // GetProcessMemoryInfo() + #if defined(MS_WIN32) && !defined(MS_WIN64) #define HANDLE_TO_PYNUM(handle) \ PyLong_FromUnsignedLong((unsigned long) handle) @@ -3080,6 +3087,61 @@ _winapi_ReportEvent_impl(PyObject *module, HANDLE handle, } +/*[clinic input] +_winapi.GetProcessMemoryInfo + handle: HANDLE + / + +Return the memory usage of the given process handle as a dict. +[clinic start generated code]*/ + +static PyObject * +_winapi_GetProcessMemoryInfo_impl(PyObject *module, HANDLE handle) +/*[clinic end generated code: output=00a5d09732e84120 input=5b90ad61cdc68d2a]*/ +{ + 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) { \ + goto error; \ + } \ + if (PyDict_SetItemString(result, #ATTR, obj) < 0) { \ + Py_DECREF(obj); \ + goto error; \ + } \ + Py_DECREF(obj); \ + } while (0) + + ADD(PageFaultCount); + ADD(PeakWorkingSetSize); + ADD(WorkingSetSize); + ADD(QuotaPeakPagedPoolUsage); + ADD(QuotaPagedPoolUsage); + ADD(QuotaPeakNonPagedPoolUsage); + ADD(QuotaNonPagedPoolUsage); + ADD(PagefileUsage); + ADD(PeakPagefileUsage); + +#undef ADD + + return result; + +error: + Py_DECREF(result); + return NULL; +} + + static PyMethodDef winapi_functions[] = { _WINAPI_CLOSEHANDLE_METHODDEF _WINAPI_CONNECTNAMEDPIPE_METHODDEF @@ -3130,6 +3192,7 @@ static PyMethodDef winapi_functions[] = { _WINAPI__MIMETYPES_READ_WINDOWS_REGISTRY_METHODDEF _WINAPI_NEEDCURRENTDIRECTORYFOREXEPATH_METHODDEF _WINAPI_COPYFILE2_METHODDEF + _WINAPI_GETPROCESSMEMORYINFO_METHODDEF {NULL, NULL} }; @@ -3226,6 +3289,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); diff --git a/Modules/clinic/_winapi.c.h b/Modules/clinic/_winapi.c.h index 00cce91dca43b1..dd9dbffaa9ac23 100644 --- a/Modules/clinic/_winapi.c.h +++ b/Modules/clinic/_winapi.c.h @@ -2331,7 +2331,34 @@ _winapi_ReportEvent(PyObject *module, PyObject *const *args, Py_ssize_t nargs) return return_value; } +PyDoc_STRVAR(_winapi_GetProcessMemoryInfo__doc__, +"GetProcessMemoryInfo($module, handle, /)\n" +"--\n" +"\n" +"Return the memory usage of the given process handle as a dict."); + +#define _WINAPI_GETPROCESSMEMORYINFO_METHODDEF \ + {"GetProcessMemoryInfo", (PyCFunction)_winapi_GetProcessMemoryInfo, METH_O, _winapi_GetProcessMemoryInfo__doc__}, + +static PyObject * +_winapi_GetProcessMemoryInfo_impl(PyObject *module, HANDLE handle); + +static PyObject * +_winapi_GetProcessMemoryInfo(PyObject *module, PyObject *arg) +{ + PyObject *return_value = NULL; + HANDLE handle; + + if (!PyArg_Parse(arg, "" F_HANDLE ":GetProcessMemoryInfo", &handle)) { + goto exit; + } + return_value = _winapi_GetProcessMemoryInfo_impl(module, handle); + +exit: + return return_value; +} + #ifndef _WINAPI_GETSHORTPATHNAME_METHODDEF #define _WINAPI_GETSHORTPATHNAME_METHODDEF #endif /* !defined(_WINAPI_GETSHORTPATHNAME_METHODDEF) */ -/*[clinic end generated code: output=4ab94eaee93a0a90 input=a9049054013a1b77]*/ +/*[clinic end generated code: output=07dfd4bbacaed4a8 input=a9049054013a1b77]*/

pFad - (p)hone/(F)rame/(a)nonymizer/(d)eclutterfier! Saves Data!