Rewrite the code to load icon resources for default programs

Drops the dependency on pywin32 and is hopefully more robust as well

src/calibre/utils/open_with/windows.py

@@ -5,23 +5,20 @@
 __license__ = 'GPL v3'
 __copyright__ = '2015, Kovid Goyal <kovid at kovidgoyal.net>'
 
-import re, struct, ctypes
+import re
-from collections import namedtuple
+from PyQt5.Qt import QBuffer, QByteArray, QPixmap, Qt, QtWin
-from polyglot.builtins import map
 
-from PyQt5.Qt import QtWin, Qt, QIcon, QByteArray, QBuffer, QPixmap
+from calibre.constants import plugins
-import win32con, win32api, win32gui, pywintypes, winerror
-
-from calibre import prints
 from calibre.gui2 import must_use_qt
 from calibre.utils.winreg.default_programs import split_commandline
 from polyglot.builtins import filter
 
 ICON_SIZE = 64
+winutil = plugins['winutil'][0]
 
 
 def hicon_to_pixmap(hicon):
-    return QtWin.fromHICON(hicon)
+    return QtWin.fromHICON(int(hicon))
 
 
 def pixmap_to_data(pixmap):
@@ -32,84 +29,6 @@ def pixmap_to_data(pixmap):
     return bytearray(ba.data())
 
 
-def copy_to_size(pixmap, size=ICON_SIZE):
-    if pixmap.width() > ICON_SIZE:
-        return pixmap.scaled(ICON_SIZE, ICON_SIZE, transformMode=Qt.SmoothTransformation)
-    return pixmap.copy()
-
-
-def simple_load_icon(module, index, as_data=False, size=ICON_SIZE):
-    ' Use the win32 API ExtractIcon to load the icon. This restricts icon size to 32x32, but has less chance of failing '
-    try:
-        large_icons, small_icons = win32gui.ExtractIconEx(module, index, 10)
-    except pywintypes.error as err:
-        if err.winerror != winerror.ERROR_FILE_NOT_FOUND:
-            raise
-        prints('File %r does not exist, cannot load icon' % module)
-        return
-    icons = large_icons + small_icons
-    try:
-        if icons:
-            must_use_qt()
-            pixmap = copy_to_size(QtWin.fromHICON(icons[0]), size=size)
-            if as_data:
-                return pixmap_to_data(pixmap)
-            return QIcon(pixmap)
-    finally:
-        tuple(map(win32gui.DestroyIcon, icons))
-
-
-def read_icon(handle, icon):
-    must_use_qt()
-    resource = win32api.LoadResource(handle, win32con.RT_ICON, icon.id)
-    pixmap = QPixmap()
-    pixmap.loadFromData(resource)
-    hicon = None
-    if pixmap.isNull():
-        if icon.width > 0 and icon.height > 0:
-            hicon = ctypes.windll.user32.CreateIconFromResourceEx(
-                resource, len(resource), True, 0x00030000, icon.width, icon.height, win32con.LR_DEFAULTCOLOR)
-        else:
-            hicon = win32gui.CreateIconFromResource(resource, True)
-        pixmap = hicon_to_pixmap(hicon).copy()
-        win32gui.DestroyIcon(hicon)
-    return pixmap
-
-
-def load_icon(module, index, as_data=False, size=ICON_SIZE):
-    handle = win32api.LoadLibraryEx(module, 0, 0x20 | 0x2)
-    try:
-        ids = win32api.EnumResourceNames(handle, win32con.RT_GROUP_ICON)
-        grp_id = -index if index < 0 else ids[index]
-        data = win32api.LoadResource(handle, win32con.RT_GROUP_ICON, grp_id)
-        pos = 0
-        reserved, idtype, count = struct.unpack_from(b'<HHH', data, pos)
-        pos += 6
-        fmt = b'<BBBBHHIH'
-        ssz = struct.calcsize(fmt)
-        icons = []
-        Icon = namedtuple('Icon', 'size width height id')
-        while count > 0:
-            count -= 1
-            width, height, color_count, reserved, planes, bitcount, isize, icon_id = struct.unpack_from(fmt, data, pos)
-            icons.append(Icon(isize, width, height, icon_id))
-            pos += ssz
-        # Unfortunately we cannot simply choose the icon with the largest
-        # width/height as the width and height are not recorded for PNG images
-        pixmaps = []
-        for icon in icons:
-            ans = read_icon(handle, icon)
-            if ans is not None and not ans.isNull():
-                pixmaps.append(ans)
-        pixmaps.sort(key=lambda p:p.size().width(), reverse=True)
-        pixmap = copy_to_size(pixmaps[0], size=size)
-        if as_data:
-            pixmap = pixmap_to_data(pixmap)
-        return pixmap
-    finally:
-        win32api.FreeLibrary(handle)
-
-
 def load_icon_resource(icon_resource, as_data=False, size=ICON_SIZE):
     if not icon_resource:
         return
@@ -120,7 +39,24 @@ def load_icon_resource(icon_resource, as_data=False, size=ICON_SIZE):
     index = int(index)
     if module.startswith('"') and module.endswith('"'):
         module = split_commandline(module)[0]
-    try:
+    hmodule = winutil.load_library(module, winutil.LOAD_LIBRARY_AS_DATAFILE | winutil.LOAD_LIBRARY_AS_IMAGE_RESOURCE)
-        return load_icon(module, index, as_data=as_data, size=size)
+    icons = winutil.load_icons(hmodule, index)
-    except Exception:
+    pixmaps = []
-        return simple_load_icon(module, index, as_data=as_data, size=size)
+    must_use_qt()
+    for icon_data, icon_handle in icons:
+        pixmap = QPixmap()
+        pixmap.loadFromData(icon_data)
+        if pixmap.isNull() and bool(icon_handle):
+            pixmap = hicon_to_pixmap(icon_handle)
+        if pixmap.isNull():
+            continue
+        pixmaps.append(pixmap)
+    if not pixmaps:
+        return
+    pixmaps.sort(key=lambda p: p.width())
+    for pmap in pixmaps:
+        if pmap.width() >= size:
+            if pmap.width() == size:
+                return pmap
+            return pixmap.scaled(size, size, transformMode=Qt.SmoothTransformation)
+    return pixmaps[-1].scaled(size, size, transformMode=Qt.SmoothTransformation)

src/calibre/utils/windows/winutil.cpp

@@ -55,6 +55,11 @@ Handle_as_int(Handle * self) {
 	return PyLong_FromVoidPtr(self->handle);
 }
 
+static int
+Handle_as_bool(Handle *self) {
+	return self->handle != NULL;
+}
+
 static PyObject*
 Handle_repr(Handle * self) {
 	const char* name = "UNKNOWN";
@@ -736,7 +741,7 @@ get_process_times(PyObject *self, PyObject *pid) {
     }
     FILETIME creation, exit, kernel, user;
     BOOL ok = GetProcessTimes(h, &creation, &exit, &kernel, &user);
-    int ec = GetLastError();
+    DWORD ec = GetLastError();
     CloseHandle(h);
     if (!ok) return PyErr_SetFromWindowsErr(ec);
 #define T(ft) ((unsigned long long)(ft.dwHighDateTime) << 32 | ft.dwLowDateTime)
@@ -763,7 +768,7 @@ get_handle_information(PyObject *self, PyObject *args) {
 
 static PyObject*
 get_last_error(PyObject *self, PyObject *args) {
-	return PyLong_FromLong(GetLastError());
+	return PyLong_FromUnsignedLong(GetLastError());
 }
 
 static PyObject*
@@ -776,6 +781,92 @@ load_library(PyObject *self, PyObject *args) {
 	return (PyObject*)Handle_create(h, ModuleHandle, PyTuple_GET_ITEM(args, 0));
 }
 
+typedef struct {
+	int count;
+	const wchar_t *resource_id;
+} ResourceData;
+
+
+BOOL CALLBACK
+EnumResProc(HMODULE handle, LPWSTR type, LPWSTR name, ResourceData *data) {
+	if (data->count-- > 0) return TRUE;
+	data->resource_id = name;
+	return FALSE;
+}
+
+static const wchar_t*
+get_resource_id_for_index(HMODULE handle, const int index, LPCWSTR type = RT_GROUP_ICON) {
+	ResourceData data = {index, NULL};
+	int count = index;
+	EnumResourceNamesW(handle, type, reinterpret_cast<ENUMRESNAMEPROC>(EnumResProc), reinterpret_cast<LONG_PTR>(&data));
+	return data.resource_id;
+}
+
+struct GRPICONDIRENTRY {
+    BYTE bWidth;
+    BYTE bHeight;
+    BYTE bColorCount;
+    BYTE bReserved;
+    WORD wPlanes;
+    WORD wBitCount;
+    DWORD dwBytesInRes;
+    WORD nID;
+  };
+
+static PyObject*
+load_icon(PyObject *args, HMODULE handle, GRPICONDIRENTRY *entry) {
+	HRSRC res = FindResourceExW(handle, RT_ICON, MAKEINTRESOURCEW(entry->nID), MAKELANGID(LANG_NEUTRAL, SUBLANG_NEUTRAL));
+	if (!res) {
+		DWORD ec = GetLastError();
+		if (ec == ERROR_RESOURCE_TYPE_NOT_FOUND || ec == ERROR_RESOURCE_NAME_NOT_FOUND || ec == ERROR_RESOURCE_LANG_NOT_FOUND) return NULL;
+		return set_error_from_handle(args, ec);
+	}
+	HGLOBAL hglob = LoadResource(handle, res);
+	if (hglob == NULL) return set_error_from_handle(args);
+	BYTE* data = (BYTE*)LockResource(hglob);
+	if (!data) return NULL;
+	DWORD sz = SizeofResource(handle, res);
+	if (!sz) return NULL;
+	HICON icon = CreateIconFromResourceEx(data, sz, TRUE, 0x00030000, entry->bWidth, entry->bHeight, LR_DEFAULTCOLOR);
+	return Py_BuildValue("y#N", data, sz, Handle_create(icon, IconHandle));
+}
+
+struct GRPICONDIR {
+    WORD idReserved;
+    WORD idType;
+    WORD idCount;
+    GRPICONDIRENTRY idEntries[1];
+};
+
+static PyObject*
+load_icons(PyObject *self, PyObject *args) {
+	HMODULE handle;
+	int index;
+	if (!PyArg_ParseTuple(args, "O&i", convert_handle, &handle, &index)) return NULL;
+
+	LPCWSTR resource_id = index < 0 ? MAKEINTRESOURCEW(-index) : get_resource_id_for_index(handle, index);
+	if (resource_id == NULL) { PyErr_Format(PyExc_IndexError, "no resource found with index: %d in handle: %S", index, PyTuple_GET_ITEM(args, 0)); return NULL; }
+	HRSRC res = FindResourceExW(handle, RT_GROUP_ICON, resource_id, MAKELANGID(LANG_NEUTRAL, SUBLANG_NEUTRAL));
+	if (res == NULL) return set_error_from_handle(args);
+	DWORD size = SizeofResource(handle, res);
+	if (!size) { PyErr_SetString(PyExc_ValueError, "the icon group resource at the specified index is empty"); return NULL; }
+	HGLOBAL hglob = LoadResource(handle, res);
+	if (hglob == NULL) return set_error_from_handle(args);
+	GRPICONDIR *grp_icon_dir = (GRPICONDIR*)LockResource(hglob);
+	if (!grp_icon_dir) { PyErr_SetString(PyExc_RuntimeError, "failed to lock icon group resource"); return NULL; }
+	PyObject *ans = PyList_New(0);
+	if (!ans) return NULL;
+	for (size_t i = 0; i < grp_icon_dir->idCount; i++) {
+		PyObject *hicon = load_icon(args, handle, grp_icon_dir->idEntries + i);
+		if (hicon) {
+			int ret = PyList_Append(ans, hicon);
+			Py_CLEAR(hicon);
+			if (ret != 0) { Py_CLEAR(ans); return NULL; }
+		} else if (PyErr_Occurred()) { Py_CLEAR(ans); return NULL; }
+	}
+	return ans;
+}
+
 // Boilerplate  {{{
 static const char winutil_doc[] = "Defines utility methods to interface with windows.";
 
@@ -790,6 +881,7 @@ static PyMethodDef winutil_methods[] = {
 	M(get_handle_information, METH_VARARGS),
 	M(get_last_error, METH_NOARGS),
 	M(load_library, METH_VARARGS),
+	M(load_icons, METH_VARARGS),
 
     {"special_folder_path", winutil_folder_path, METH_VARARGS,
     "special_folder_path(csidl_id) -> path\n\n"
@@ -930,6 +1022,7 @@ extern "C" {
 
 CALIBRE_MODINIT_FUNC PyInit_winutil(void) {
 	HandleNumberMethods.nb_int = (unaryfunc)Handle_as_int;
+	HandleNumberMethods.nb_bool = (inquiry)Handle_as_bool;
     HandleType.tp_name = "winutil.Handle";
     HandleType.tp_doc = "Wrappers for Win32 handles that free the handle on delete automatically";
     HandleType.tp_basicsize = sizeof(Handle);