Exposed set_icon in _macosx, wrapped for other backends, used in Qt5

joelfrederico · joelfrederico · commit 994e7d661cc1 · 2019-07-31T23:18:43.000-07:00
diff --git a/lib/matplotlib/backends/backend_macos_common.py b/lib/matplotlib/backends/backend_macos_common.py
@@ -0,0 +1,9 @@
+import sys
+
+if sys.platform != "darwin":
+    def set_mac_icon(path):
+        pass
+else:
+    from . import _macosx
+    def set_mac_icon(path):
+        _macosx.set_icon(path)
diff --git a/lib/matplotlib/backends/backend_qt5.py b/lib/matplotlib/backends/backend_qt5.py
@@ -540,6 +540,12 @@ def __init__(self, canvas, num):
 
         self.window.setWindowTitle("Figure %d" % num)
         image = str(cbook._get_data_path('images/matplotlib.svg'))
+
+        if sys.platform == "darwin":
+            from . import backend_macos_common
+            image = str(cbook._get_data_path('images/matplotlib.pdf'))
+            backend_macos_common.set_mac_icon(image)
+
         self.window.setWindowIcon(QtGui.QIcon(image))
 
         # Give the keyboard focus to the figure instead of the
diff --git a/src/_macosx.m b/src/_macosx.m
@@ -178,76 +178,129 @@ static int wait_for_stdin(void)
     return 1;
 }
 
-static int set_icon(void)
+PyObject* _get_pdf_icon_path(void)
 {
     PyObject* mpl = PyImport_ImportModule("matplotlib");
     if (!mpl)
     {
-      return -1;
+        return NULL;
     }
 
     PyObject* get_data_path = PyObject_GetAttrString(mpl, "get_data_path");
     if (!get_data_path)
     {
-      Py_DECREF(mpl);
-      return -1;
+        Py_DECREF(mpl);
+        return NULL;
     }
 
     PyObject* arg = PyTuple_New(0);
     if (!arg)
     {
-      Py_DECREF(mpl);
-      Py_DECREF(get_data_path);
-      return -1;
+        Py_DECREF(mpl);
+        Py_DECREF(get_data_path);
+        return NULL;
     }
 
     PyObject* path = PyObject_Call(get_data_path, arg, NULL);
     if (!path)
     {
-      Py_DECREF(mpl);
-      Py_DECREF(get_data_path);
-      Py_DECREF(arg);
-      return -1;
+        Py_DECREF(mpl);
+        Py_DECREF(get_data_path);
+        Py_DECREF(arg);
+        return NULL;
     }
 
     PyObject* pdf_rel = PyUnicode_FromString("/images/matplotlib.pdf");
+    if (!pdf_rel)
+    {
+        Py_DECREF(mpl);
+        Py_DECREF(get_data_path);
+        Py_DECREF(arg);
+        Py_DECREF(path);
+        return NULL;
+    }
+
     PyObject* pdf_absolute = PyUnicode_Concat(path, pdf_rel);
-    if (!pdf_absolute)
+
+    Py_DECREF(mpl);
+    Py_DECREF(get_data_path);
+    Py_DECREF(arg);
+    Py_DECREF(path);
+    Py_DECREF(pdf_rel);
+
+    return pdf_absolute;
+}
+
+static PyObject* set_icon(PyObject* ignored, PyObject* path)
+{
+    if (!path)
     {
-      Py_DECREF(mpl);
-      Py_DECREF(get_data_path);
-      Py_DECREF(arg);
-      Py_DECREF(path);
-      Py_DECREF(pdf_rel);
+        PyErr_SetString(PyExc_RuntimeError, "No path passed");
+        return NULL;
     }
 
     Py_ssize_t size;
-    const char* path_string = PyUnicode_AsUTF8AndSize(pdf_absolute, &size);
+    // API indicates Python owns path_string; we don't need to free it
+    const char* path_string = PyUnicode_AsUTF8AndSize(path, &size);
     if (!path_string)
     {
-      Py_DECREF(mpl);
-      Py_DECREF(get_data_path);
-      Py_DECREF(arg);
-      Py_DECREF(path);
-      Py_DECREF(pdf_rel);
-      Py_DECREF(pdf_absolute);
-      return -1;
+        PyErr_SetString(PyExc_TypeError, "bad argument type, must be UTF-8 "
+                "path to icon");
+        return NULL;
     }
 
-    NSString* path_nsstring =
-      [[NSString alloc] initWithUTF8String: path_string];
+    NSAutoreleasePool* pool = [[NSAutoreleasePool alloc] init];
 
-    NSImage* image = [[NSImage alloc] initByReferencingFile: path_nsstring];
+    NSImage* image;
+    @try
+    {
+        NSString* path_nsstring =
+            [[NSString alloc] initWithUTF8String: path_string];
 
-    NSApplication* app = [NSApplication sharedApplication];
-    app.applicationIconImage = image;
+        image =
+            [[NSImage alloc] initByReferencingFile: path_nsstring];
+    }
+    @catch (NSException *e)
+    {
 
-    Py_DECREF(mpl);
-    Py_DECREF(get_data_path);
-    Py_DECREF(arg);
-    Py_DECREF(path);
-    Py_DECREF(pdf_rel);
-    Py_DECREF(pdf_absolute);
+        PyErr_SetString(PyExc_RuntimeError, "could not load icon");
+        [pool drain];
+        return NULL;
+    }
+
+    @try
+    {
+        NSApplication* app = [NSApplication sharedApplication];
+        app.applicationIconImage = image;
+    }
+    @catch (NSException *e)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "could not set icon");
+        [pool drain];
+        return NULL;
+    }
+
+    [pool drain];
+    Py_RETURN_NONE;
+}
+
+static int _set_icon(void)
+{
+    PyObject* pdf_path = _get_pdf_icon_path();
+    if (!pdf_path)
+    {
+        return -1;
+    }
+
+    PyObject* result = set_icon(NULL, pdf_path);
+    if (!result)
+    {
+        Py_DECREF(pdf_path);
+        return -1;
+    }
+
+    Py_DECREF(pdf_path);
+    Py_DECREF(result);
     return 0;
 }
 
@@ -813,7 +866,10 @@ static CGFloat _get_device_scale(CGContextRef cr)
 
     NSAutoreleasePool* pool = [[NSAutoreleasePool alloc] init];
 
-    set_icon();
+    if (_set_icon() != 0)
+    {
+        return -1;
+    }
 
     self->window = [self->window initWithContentRect: rect
                                          styleMask: NSTitledWindowMask
@@ -2702,6 +2758,11 @@ static bool verify_framework(void)
     METH_VARARGS,
     "Sets the active cursor."
    },
+   {"set_icon",
+    set_icon,
+    METH_O,
+    "Sets the Dock icon in macOS"
+   },
    {NULL, NULL, 0, NULL} /* sentinel */
 };
 
