diff --git a/comtypes/malloc.py b/comtypes/malloc.py new file mode 100644 index 00000000..d4543d41 --- /dev/null +++ b/comtypes/malloc.py @@ -0,0 +1,77 @@ +import logging +from ctypes import ( + HRESULT, + POINTER, + OleDLL, + WinDLL, + byref, + c_int, + c_size_t, + c_ulong, + c_void_p, + c_wchar, + c_wchar_p, + cast, + memmove, + sizeof, + wstring_at, +) +from ctypes.wintypes import DWORD, LPVOID + +from comtypes import COMMETHOD, GUID, IUnknown +from comtypes.GUID import _CoTaskMemFree + +logger = logging.getLogger(__name__) + + +class IMalloc(IUnknown): + _iid_ = GUID("{00000002-0000-0000-C000-000000000046}") + _methods_ = [ + COMMETHOD([], c_void_p, "Alloc", ([], c_ulong, "cb")), + COMMETHOD([], c_void_p, "Realloc", ([], c_void_p, "pv"), ([], c_ulong, "cb")), + COMMETHOD([], None, "Free", ([], c_void_p, "py")), + COMMETHOD([], c_ulong, "GetSize", ([], c_void_p, "pv")), + COMMETHOD([], c_int, "DidAlloc", ([], c_void_p, "pv")), + COMMETHOD([], None, "HeapMinimize"), # 25 + ] + + +_ole32 = OleDLL("ole32") + +_CoGetMalloc = _ole32.CoGetMalloc +_CoGetMalloc.argtypes = [DWORD, POINTER(POINTER(IMalloc))] +_CoGetMalloc.restype = HRESULT + +_ole32_nohresult = WinDLL("ole32") + +SIZE_T = c_size_t +_CoTaskMemAlloc = _ole32_nohresult.CoTaskMemAlloc +_CoTaskMemAlloc.argtypes = [SIZE_T] +_CoTaskMemAlloc.restype = LPVOID + +malloc = POINTER(IMalloc)() +_CoGetMalloc(1, byref(malloc)) +assert bool(malloc) + + +def from_outparam(self): + if not self: + return None + result = wstring_at(self) + # `DidAlloc` method returns; + # * 1 (allocated) + # * 0 (not allocated) + # * -1 (cannot determine or NULL) + # https://learn.microsoft.com/en-us/windows/win32/api/objidl/nf-objidl-imalloc-didalloc + assert malloc.DidAlloc(self), "memory was NOT allocated by CoTaskMemAlloc" + _CoTaskMemFree(self) + return result + + +def comstring(text, typ=c_wchar_p): + size = (len(text) + 1) * sizeof(c_wchar) + mem = _CoTaskMemAlloc(size) + logger.debug("malloc'd 0x%x, %d bytes" % (mem, size)) + ptr = cast(mem, typ) + memmove(mem, text, size) + return ptr diff --git a/comtypes/test/test_outparam.py b/comtypes/test/test_from_outparam.py similarity index 100% rename from comtypes/test/test_outparam.py rename to comtypes/test/test_from_outparam.py diff --git a/comtypes/test/test_malloc.py b/comtypes/test/test_malloc.py new file mode 100644 index 00000000..597e58b7 --- /dev/null +++ b/comtypes/test/test_malloc.py @@ -0,0 +1,6 @@ +import unittest as ut + +from comtypes.malloc import IMalloc # noqa + + +class Test(ut.TestCase): ... diff --git a/pyproject.toml b/pyproject.toml index f9cea138..b44e4dd3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,7 +89,7 @@ ignore = ["E402"] "comtypes/test/test_client.py" = ["F401"] "comtypes/test/test_dict.py" = ["F841"] "comtypes/test/test_eventinterface.py" = ["F841"] -"comtypes/test/test_outparam.py" = ["F841"] +"comtypes/test/test_from_outparam.py" = ["F841"] "comtypes/test/test_sapi.py" = ["E401"] "comtypes/test/test_server.py" = ["F401", "F841"] "comtypes/test/test_subinterface.py" = ["E401", "F401", "F403", "F405"]