More work on the file dialogs helper

setup/installer/windows/file_dialogs.cpp

@@ -5,22 +5,110 @@
  * Distributed under terms of the GPL3 license.
  */
 
+#ifndef _UNICODE
+#define _UNICODE
+#endif
 #include <Windows.h>
 #include <Shobjidl.h>
+#include <comdef.h>
 #include <stdio.h>
 #include <string.h>
 #include <io.h>
 #include <fcntl.h>
 
-#define PRINTERR(x) fprintf(stderr, x); fflush(stderr);
+#define PRINTERR(x) fprintf(stderr, "%s", x); fflush(stderr);
-#define REPORTERR(x) { PRINTERR(x); ret = 1; goto error; }
-#define CALLCOM(x, err) hr = x; if(FAILED(hr)) REPORTERR(err)
 
-int show_dialog(HWND parent, bool save_dialog, LPWSTR title, LPWSTR folder, LPWSTR filename, LPWSTR save_path, bool multiselect, bool confirm_overwrite, bool only_dirs, bool no_symlinks) {
+bool write_bytes(size_t sz, const char* buf) {
-	int ret = 0;
+    size_t num = 0;
+    while(sz > 0 && !feof(stdout) && !ferror(stdout)) {
+        num = fwrite(buf, sizeof(char), sz, stdout);
+        if (num == 0) break;
+        buf += num; sz -= num;
+    }
+    if (sz > 0) PRINTERR("Failed to write to stdout");
+    return sz == 0;
+}
+
+bool read_bytes(size_t sz, char* buf, bool allow_incomplete=false) {
+	char *ptr = buf, *limit = buf + sz;
+	while(limit > ptr && !feof(stdin) && !ferror(stdin)) {
+		ptr += fread(ptr, 1, limit - ptr, stdin);
+	}
+	if (ferror(stdin)) { PRINTERR("Failed to read from stdin!"); return false; }
+	if (ptr - buf != sz) { if (!allow_incomplete) PRINTERR("Truncated input!"); return false; }
+	return true;
+}
+
+bool from_utf8(size_t sz, const char *src, LPWSTR* ans) {
+	int asz = MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, src, (int)sz, NULL, 0);
+	if (!asz) { PRINTERR("Failed to get size of UTF-8 string"); return false; }
+	*ans = (LPWSTR)calloc(asz+1, sizeof(wchar_t));
+	if(*ans == NULL) { PRINTERR("Out of memory!"); return false; }
+	asz = MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, src, (int)sz, *ans, asz);
+	if (!asz) { PRINTERR("Failed to convert UTF-8 string"); return false; }
+	return true;
+}
+
+char* to_utf8(LPCWSTR src, int *sz) {
+    char *ans = NULL;
+    *sz = WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS, src, -1, NULL, 0, NULL, NULL);
+	if (!*sz) { PRINTERR("Failed to get size of UTF-16 string"); return NULL; }
+    ans = (char*)calloc((*sz) + 1, sizeof(char));
+    if (ans == NULL) { PRINTERR("Out of memory!"); return NULL; }
+    *sz = WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS, src, -1, ans, *sz, NULL, NULL);
+    if (!*sz) { PRINTERR("Failed to convert UTF-16 string"); return NULL; }
+    return ans;
+}
+
+static char* rsbuf = NULL;
+
+bool read_string(unsigned short sz, LPWSTR* ans) {
+	memset(rsbuf, 0, 65537);
+	if (!read_bytes(sz, rsbuf)) return false;
+	if (!from_utf8(sz, rsbuf, ans)) return false;
+	return true;
+}
+
+COMDLG_FILTERSPEC *read_file_types(UINT *num_file_types) {
+    char buf[10] = {0};
+    COMDLG_FILTERSPEC *ans = NULL;
+
+    if(!read_bytes(sizeof(unsigned short), buf)) return NULL;
+    *num_file_types = *((unsigned short*)buf);
+    if (*num_file_types < 1 || *num_file_types > 500) { PRINTERR("Invalid number of file types"); return NULL; }
+    ans = (COMDLG_FILTERSPEC*)calloc((*num_file_types) + 1, sizeof(COMDLG_FILTERSPEC));
+    if (ans == NULL) { PRINTERR("Out of memory!"); return NULL; }
+
+    for(unsigned short i = 0; i < *num_file_types; i++) {
+        if(!read_bytes(sizeof(unsigned short), buf)) return NULL;
+        if(!read_string(*((unsigned short*)buf), (LPWSTR*)&(ans[i].pszName))) return NULL;
+        if(!read_bytes(sizeof(unsigned short), buf)) return NULL;
+        if(!read_string(*((unsigned short*)buf), (LPWSTR*)&(ans[i].pszSpec))) return NULL;
+    }
+    return ans;
+}
+
+static void print_com_error(HRESULT hr, const char *msg) {
+    _com_error err(hr);
+    LPCWSTR emsg = (LPCWSTR) err.ErrorMessage();
+    int sz = 0;
+    const char *buf = to_utf8(emsg, &sz);
+    if (buf == NULL) { fprintf(stderr, "%s", msg); }
+    else { fprintf(stderr, "%s: (HRESULT=0x%x) %s\n", msg, hr, emsg); }
+    fflush(stderr);
+}
+
+#define REPORTERR(hr, x) { print_com_error(hr, x); ret = 1; goto error; }
+#define CALLCOM(x, err) hr = x; if(FAILED(hr)) REPORTERR(hr, err)
+
+int show_dialog(HWND parent, bool save_dialog, LPWSTR title, LPWSTR folder, LPWSTR filename, LPWSTR save_path, bool multiselect, bool confirm_overwrite, bool only_dirs, bool no_symlinks, COMDLG_FILTERSPEC *file_types, UINT num_file_types) {
+	int ret = 0, name_sz = 0;
 	IFileDialog *pfd = NULL;
-	IShellItem *result = NULL, *folder_item = NULL, *save_path_item = NULL;
+	IShellItemArray *items = NULL;
-	DWORD options;
+    IShellItem *item = NULL, *folder_item = NULL, *save_path_item = NULL;
+    char *path = NULL;
+	DWORD options = 0, item_count = 0;
+    LPWSTR name = NULL;
 	HRESULT hr = S_OK;
 	hr = CoInitialize(NULL);
 	if (FAILED(hr)) { PRINTERR("Failed to initialize COM"); return 1; }
@@ -53,51 +141,41 @@ int show_dialog(HWND parent, bool save_dialog, LPWSTR title, LPWSTR folder, LPWS
 		if (SUCCEEDED(hr)) pfd->SetFolder(folder_item);
 	}
 	if (filename != NULL) pfd->SetFileName(filename); // Failure is not critical
+    if (!(options & FOS_PICKFOLDERS) && file_types != NULL && num_file_types > 0) {
+        CALLCOM(pfd->SetFileTypes(num_file_types, file_types), "Failed to set file types")
+    }
 	hr = pfd->Show(parent);
 	if (hr == HRESULT_FROM_WIN32(ERROR_CANCELLED)) goto error;
-	if (FAILED(hr)) REPORTERR("Failed to show dialog")
+	if (FAILED(hr)) REPORTERR(hr, "Failed to show dialog")
 
-	CALLCOM(pfd->GetResult(&result), "Failed to get dialog result")
+    if (save_dialog) {
+        CALLCOM(pfd->GetResult(&item), "Failed to get save dialog result");
+        CALLCOM(item->GetDisplayName(SIGDN_FILESYSPATH, &name), "Failed to get display name of save dialog result");
+        path = to_utf8(name, &name_sz);
+        CoTaskMemFree(name); name = NULL;
+        if (path == NULL) return 1;
+        if (!write_bytes(name_sz, path)) return 1;
+    } else {
+        CALLCOM(((IFileOpenDialog*)pfd)->GetResults(&items), "Failed to get dialog results");
+        CALLCOM(items->GetCount(&item_count), "Failed to get count of results");
+        if (item_count > 0) {
+            for (DWORD i = 0; i < item_count; i++) {
+                CALLCOM(items->GetItemAt(i, &item), "Failed to get result item");
+                if (SUCCEEDED(item->GetDisplayName(SIGDN_FILESYSPATH, &name))) {
+                    path = to_utf8(name, &name_sz);
+                    CoTaskMemFree(name); name = NULL;
+                    if (path == NULL) return 1;
+                    if (!write_bytes(name_sz, path)) return 1;
+                }
+            }
+        }
+    }
 
 error:
 	if(pfd) pfd->Release();
 	CoUninitialize();
 	return ret;
 }
-
-bool read_bytes(size_t sz, char* buf, bool allow_incomplete=false) {
-	char *ptr = buf, *limit = buf + sz;
-	while(limit > ptr && !feof(stdin) && !ferror(stdin)) {
-		ptr += fread(ptr, 1, limit - ptr, stdin);
-	}
-	if (ferror(stdin)) { PRINTERR("Failed to read from stdin!"); return false; }
-	if (ptr - buf != sz) { if (!allow_incomplete) PRINTERR("Truncated input!"); return false; }
-	return true;
-}
-
-bool from_utf8(size_t sz, const char *src, LPWSTR* ans) {
-	int asz = MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, src, (int)sz, NULL, 0);
-	if (!asz) { PRINTERR("Failed to get size of UTF-8 string"); return false; }
-	*ans = (LPWSTR)calloc(asz+1, sizeof(wchar_t));
-	if(*ans == NULL) { PRINTERR("Out of memory!"); return false; }
-	asz = MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, src, (int)sz, *ans, asz);
-	if (!asz) { PRINTERR("Failed to convert UTF-8 string"); return false; }
-	return true;
-}
-
-static char* rsbuf = NULL;
-
-bool read_string(unsigned short sz, LPWSTR* ans) {
-	if(rsbuf == NULL) {
-		rsbuf = (char*)calloc(65537, sizeof(char));
-		if(rsbuf == NULL) { PRINTERR("Out of memory!"); return false; }
-	}
-	memset(rsbuf, 0, 65537);
-	if (!read_bytes(sz, rsbuf)) return false;
-	if (!from_utf8(sz, rsbuf, ans)) return false;
-	return true;
-}
-
 #define READ(x, y) if (!read_bytes((x), (y))) return 1;
 #define CHECK_KEY(x) (key_size == sizeof(x) - 1 && memcmp(buf, x, sizeof(x) - 1) == 0)
 #define READSTR(x) READ(sizeof(unsigned short), buf); if(!read_string(*((unsigned short*)buf), &x)) return 1;
@@ -105,17 +183,21 @@ bool read_string(unsigned short sz, LPWSTR* ans) {
 #define READBOOL(x)  READ(1, buf); x = !!buf[0]; 
 
 int WINAPI wWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, PWSTR pCmdLine, int nCmdShow) {
-	char buf[257];
+	char buf[257] = {0};
 	size_t key_size = 0;
 	HWND parent = NULL;
 	bool save_dialog = false, multiselect = false, confirm_overwrite = false, only_dirs = false, no_symlinks = false;
 	unsigned short len = 0;
 	LPWSTR title = NULL, folder = NULL, filename = NULL, save_path = NULL;
+    COMDLG_FILTERSPEC *file_types = NULL;
+    UINT num_file_types = 0;
 
 	SETBINARY(stdout); SETBINARY(stdin); SETBINARY(stderr);
 	// The calibre executables call SetDllDirectory, we unset it here just in
 	// case it interferes with some idiotic shell extension or the other
 	SetDllDirectory(NULL);
+    rsbuf = (char*)calloc(65537, sizeof(char));
+    if(rsbuf == NULL) { PRINTERR("Out of memory!"); return 1; }
 
 	while(!feof(stdin)) {
 		memset(buf, 0, sizeof(buf));
@@ -147,11 +229,13 @@ int WINAPI wWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, PWSTR pCmdLine
 
 		else if CHECK_KEY("NO_SYMLINKS") { READBOOL(no_symlinks) }
 
+        else if CHECK_KEY("FILE_TYPES") { file_types = read_file_types(&num_file_types); if (file_types == NULL) return 1; }
+
 		else {
 			PRINTERR("Unknown key");
 			return 1;
 		}
 	}
 
-	return show_dialog(parent, save_dialog, title, folder, filename, save_path, multiselect, confirm_overwrite, only_dirs, no_symlinks);
+	return show_dialog(parent, save_dialog, title, folder, filename, save_path, multiselect, confirm_overwrite, only_dirs, no_symlinks, file_types, num_file_types);
 }

src/calibre/gui2/win_file_dialogs.py

@@ -44,6 +44,17 @@ def serialize_string(key, val):
         raise ValueError('%s is too long' % key)
     return struct.pack(b'=B%dsH%ds' % (len(key), len(val)), len(key), key, len(val), val)
 
+def serialize_file_types(file_types):
+    key = b"FILE_TYPES"
+    buf = [struct.pack(b'=B%dsH' % len(key), len(key), key, len(file_types))]
+    def add(x):
+        x = x.encode('utf-8').replace(b'\0', b'')
+        buf.append(struct.pack(b'=H%ds' % len(x), len(x), x))
+    for name, extensions in file_types:
+        add(name or _('Files'))
+        add('; '.join('*.' + ext.lower() for ext in extensions))
+    return b''.join(buf)
+
 class Helper(Thread):
 
     def __init__(self, process, data, callback):
@@ -85,7 +96,8 @@ def select_initial_dir(q):
 
 def run_file_dialog(
         parent=None, title=None, initial_folder=None, filename=None, save_path=None,
-        allow_multiples=False, only_dirs=False, confirm_overwrite=True, save_as=False, no_symlinks=False
+        allow_multiple=False, only_dirs=False, confirm_overwrite=True, save_as=False, no_symlinks=False,
+        file_types=()
 ):
     data = []
     if parent is not None:
@@ -108,8 +120,8 @@ def run_file_dialog(
                 if not filename:
                     filename = os.path.basename(save_path)
     else:
-        if allow_multiples:
+        if allow_multiple:
-            data.append(serialize_binary('MULTISELECT', allow_multiples))
+            data.append(serialize_binary('MULTISELECT', allow_multiple))
         if only_dirs:
             data.append(serialize_binary('ONLY_DIRS', only_dirs))
     if initial_folder is not None:
@@ -120,6 +132,12 @@ def run_file_dialog(
         if isinstance(filename, bytes):
             filename = filename.decode(filesystem_encoding)
         data.append(serialize_string('FILENAME', filename))
+    if only_dirs:
+        file_types = ()  # file types not allowed for dir only dialogs
+    elif not file_types:
+        file_types = [(_('All files'), ('*',))]
+    if file_types:
+        data.append(serialize_file_types(file_types))
     loop = Loop()
     h = Helper(subprocess.Popen(
         [HELPER], stdout=subprocess.PIPE, stdin=subprocess.PIPE, stderr=subprocess.PIPE),
@@ -130,15 +148,22 @@ def run_file_dialog(
         raise Exception('File dialog failed: ' + h.stderrdata.decode('utf-8'))
     if not h.stdoutdata:
         return ()
-    return tuple(x.decode('utf-8') for x in h.stdoutdata.split(b'\0'))
+    ans = tuple(filter(None, (os.path.abspath(x.decode('utf-8')) for x in h.stdoutdata.split(b'\0'))))
+    if len(ans) > 1:
+        ans = ans[:-1]  # For some reason windows returns the initial folder as well, as the last item
+    return ans
 
 if __name__ == '__main__':
     HELPER = sys.argv[-1]
     app = QApplication([])
     q = QMainWindow()
+    _ = lambda x: x
 
     def clicked():
-        print(run_file_dialog(b, 'Testing dialogs', save_as=True, save_path='~/xxx.fdgdfg')), sys.stdout.flush()
+        print(run_file_dialog(
+            b, 'Testing dialogs', only_dirs=False, allow_multiple=True, initial_folder=expanduser('~/build/calibre'),
+            file_types=[('YAML files', ['yaml']), ('All files', '*')]))
+        sys.stdout.flush()
 
     b = QPushButton('click me')
     b.clicked.connect(clicked)