FIX: destroy figs in python when destroyed in js

tacaswell · tacaswell · commit 788894229694 · 2015-08-16T00:03:38.000-04:00
- only actually closes the figure (in the mpl sense) if all of the managers active comms are dead - closes figure when all visible copies are removed from dom - eliminates many of the [IPKernelApp] ERROR | No such comm: XXXX errors - may have caused a regression in figures not closing on kernel restart closes #4841
diff --git a/lib/matplotlib/backends/backend_nbagg.py b/lib/matplotlib/backends/backend_nbagg.py
@@ -160,16 +160,22 @@ def _create_comm(self):
 
     def destroy(self):
         self._send_event('close')
-        for comm in self.web_sockets.copy():
+        # need to copy comms as callbacks will modify this list
+        for comm in list(self.web_sockets):
             comm.on_close()
+        self.clearup_closed()
 
     def clearup_closed(self):
         """Clear up any closed Comms."""
         self.web_sockets = set([socket for socket in self.web_sockets
                                 if socket.is_open()])
 
         if len(self.web_sockets) == 0:
-            self.manager.canvas.close_event()
+            self.canvas.close_event()
+
+    def remove_comm(self, comm_id):
+        self.web_sockets = set([socket for socket in self.web_sockets
+                                if not socket.comm.comm_id == comm_id])
 
 
 class TimerTornado(TimerBase):
@@ -267,15 +273,27 @@ def __init__(self, manager):
         self.comm.on_msg(self.on_message)
 
         manager = self.manager
-        self.comm.on_close(lambda close_message: manager.clearup_closed())
+        self._ext_close = False
+
+        def _on_close(close_message):
+            self._ext_close = True
+            manager.remove_comm(close_message['content']['comm_id'])
+            manager.clearup_closed()
+
+        self.comm.on_close(_on_close)
 
     def is_open(self):
-        return not self.comm._closed
+        return not (self._ext_close or self.comm._closed)
 
     def on_close(self):
         # When the socket is closed, deregister the websocket with
         # the FigureManager.
-        self.comm.close()
+        if self.is_open():
+            try:
+                self.comm.close()
+            except KeyError:
+                # apparently already cleaned it up?
+                pass
 
     def send_json(self, content):
         self.comm.send({'data': json.dumps(content)})
@@ -298,6 +316,7 @@ def on_message(self, message):
         message = json.loads(message['content']['data'])
         if message['type'] == 'closing':
             self.on_close()
+            self.manager.clearup_closed()
         elif message['type'] == 'supports_binary':
             self.supports_binary = message['value']
         else:
diff --git a/lib/matplotlib/backends/web_backend/nbagg_mpl.js b/lib/matplotlib/backends/web_backend/nbagg_mpl.js
@@ -55,15 +55,21 @@ mpl.mpl_figure_comm = function(comm, msg) {
 };
 
 mpl.figure.prototype.handle_close = function(fig, msg) {
+    fig.root.unbind('remove')
+
     // Update the output cell to use the data from the current canvas.
     fig.push_to_output();
     var dataURL = fig.canvas.toDataURL();
     // Re-enable the keyboard manager in IPython - without this line, in FF,
     // the notebook keyboard shortcuts fail.
     IPython.keyboard_manager.enable()
     $(fig.parent_element).html('<img src="' + dataURL + '">');
-    fig.send_message('closing', {});
-    fig.ws.close()
+    fig.close_ws(fig, msg);
+}
+
+mpl.figure.prototype.close_ws = function(fig, msg){
+    fig.send_message('closing', msg);
+    // fig.ws.close()
 }
 
 mpl.figure.prototype.push_to_output = function(remove_interactive) {
@@ -126,6 +132,12 @@ mpl.figure.prototype._init_toolbar = function() {
     titlebar.prepend(buttongrp);
 }
 
+mpl.figure.prototype._root_extra_style = function(el){
+    var fig = this
+    el.on("remove", function(){
+	fig.close_ws(fig, {});
+    });
+}
 
 mpl.figure.prototype._canvas_extra_style = function(el){
     // this is important to make the div 'focusable
