Merge pull request #6216 from SylvainCorlay/comm-unregister

Unregistering comms in Comm Manager

IPython/html/static/services/kernels/js/comm.js

@@ -1,190 +1,190 b''
 // Copyright (c) IPython Development Team.
 // Distributed under the terms of the Modified BSD License.
 
 define([
     'base/js/namespace',
     'jquery',
     'base/js/utils',
 ], function(IPython, $, utils) {
     "use strict";
 
     //-----------------------------------------------------------------------
     // CommManager class
     //-----------------------------------------------------------------------
     
     var CommManager = function (kernel) {
         this.comms = {};
         this.targets = {};
         if (kernel !== undefined) {
             this.init_kernel(kernel);
         }
     };
     
     CommManager.prototype.init_kernel = function (kernel) {
         // connect the kernel, and register message handlers
         this.kernel = kernel;
         var msg_types = ['comm_open', 'comm_msg', 'comm_close'];
         for (var i = 0; i < msg_types.length; i++) {
             var msg_type = msg_types[i];
             kernel.register_iopub_handler(msg_type, $.proxy(this[msg_type], this));
         }
     };
     
     CommManager.prototype.new_comm = function (target_name, data, callbacks, metadata) {
         // Create a new Comm, register it, and open its Kernel-side counterpart
         // Mimics the auto-registration in `Comm.__init__` in the IPython Comm
         var comm = new Comm(target_name);
         this.register_comm(comm);
         comm.open(data, callbacks, metadata);
         return comm;
     };
     
     CommManager.prototype.register_target = function (target_name, f) {
         // Register a target function for a given target name
         this.targets[target_name] = f;
     };
     
     CommManager.prototype.unregister_target = function (target_name, f) {
         // Unregister a target function for a given target name
         delete this.targets[target_name];
     };
     
     CommManager.prototype.register_comm = function (comm) {
         // Register a comm in the mapping
         this.comms[comm.comm_id] = comm;
         comm.kernel = this.kernel;
         return comm.comm_id;
     };
     
-    CommManager.prototype.unregister_comm = function (comm_id) {
+    CommManager.prototype.unregister_comm = function (comm) {
         // Remove a comm from the mapping
-        delete this.comms[comm_id];
+        delete this.comms[comm.comm_id];
     };
     
     // comm message handlers
     
     CommManager.prototype.comm_open = function (msg) {
         var content = msg.content;
         var f = this.targets[content.target_name];
         if (f === undefined) {
             console.log("No such target registered: ", content.target_name);
             console.log("Available targets are: ", this.targets);
             return;
         }
         var comm = new Comm(content.target_name, content.comm_id);
         this.register_comm(comm);
         try {
             f(comm, msg);
         } catch (e) {
             console.log("Exception opening new comm:", e, e.stack, msg);
             comm.close();
             this.unregister_comm(comm);
         }
     };
     
     CommManager.prototype.comm_close = function (msg) {
         var content = msg.content;
         var comm = this.comms[content.comm_id];
         if (comm === undefined) {
             return;
         }
-        delete this.comms[content.comm_id];
+        this.unregister_comm(comm);
         try {
             comm.handle_close(msg);
         } catch (e) {
             console.log("Exception closing comm: ", e, e.stack, msg);
         }
     };
     
     CommManager.prototype.comm_msg = function (msg) {
         var content = msg.content;
         var comm = this.comms[content.comm_id];
         if (comm === undefined) {
             return;
         }
         try {
             comm.handle_msg(msg);
         } catch (e) {
             console.log("Exception handling comm msg: ", e, e.stack, msg);
         }
     };
     
     //-----------------------------------------------------------------------
     // Comm base class
     //-----------------------------------------------------------------------
     
     var Comm = function (target_name, comm_id) {
         this.target_name = target_name;
         this.comm_id = comm_id || utils.uuid();
         this._msg_callback = this._close_callback = null;
     };
     
     // methods for sending messages
     Comm.prototype.open = function (data, callbacks, metadata) {
         var content = {
             comm_id : this.comm_id,
             target_name : this.target_name,
             data : data || {},
         };
         return this.kernel.send_shell_message("comm_open", content, callbacks, metadata);
     };
     
     Comm.prototype.send = function (data, callbacks, metadata) {
         var content = {
             comm_id : this.comm_id,
             data : data || {},
         };
         return this.kernel.send_shell_message("comm_msg", content, callbacks, metadata);
     };
     
     Comm.prototype.close = function (data, callbacks, metadata) {
         var content = {
             comm_id : this.comm_id,
             data : data || {},
         };
         return this.kernel.send_shell_message("comm_close", content, callbacks, metadata);
     };
     
     // methods for registering callbacks for incoming messages
     Comm.prototype._register_callback = function (key, callback) {
         this['_' + key + '_callback'] = callback;
     };
     
     Comm.prototype.on_msg = function (callback) {
         this._register_callback('msg', callback);
     };
     
     Comm.prototype.on_close = function (callback) {
         this._register_callback('close', callback);
     };
     
     // methods for handling incoming messages
     
     Comm.prototype._maybe_callback = function (key, msg) {
         var callback = this['_' + key + '_callback'];
         if (callback) {
             try {
                 callback(msg);
             } catch (e) {
                 console.log("Exception in Comm callback", e, e.stack, msg);
             }
         }
     };
     
     Comm.prototype.handle_msg = function (msg) {
         this._maybe_callback('msg', msg);
     };
     
     Comm.prototype.handle_close = function (msg) {
         this._maybe_callback('close', msg);
     };
     
     // For backwards compatability.
     IPython.CommManager = CommManager;
     IPython.Comm = Comm;
 
     return {
         'CommManager': CommManager,
         'Comm': Comm
     };
 });

IPython/html/widgets/widget.py

@@ -1,453 +1,449 b''
 """Base Widget class.  Allows user to create widgets in the back-end that render
 in the IPython notebook front-end.
 """
 #-----------------------------------------------------------------------------
 # Copyright (c) 2013, the IPython Development Team.
 #
 # Distributed under the terms of the Modified BSD License.
 #
 # The full license is in the file COPYING.txt, distributed with this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 from contextlib import contextmanager
 
 from IPython.core.getipython import get_ipython
 from IPython.kernel.comm import Comm
 from IPython.config import LoggingConfigurable
 from IPython.utils.traitlets import Unicode, Dict, Instance, Bool, List, Tuple, Int
 from IPython.utils.py3compat import string_types
 
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------
 class CallbackDispatcher(LoggingConfigurable):
     """A structure for registering and running callbacks"""
     callbacks = List()
     
     def __call__(self, *args, **kwargs):
         """Call all of the registered callbacks."""
         value = None
         for callback in self.callbacks:
             try:
                 local_value = callback(*args, **kwargs)
             except Exception as e:
                 ip = get_ipython()
                 if ip is None:
                     self.log.warn("Exception in callback %s: %s", callback, e, exc_info=True)
                 else:
                     ip.showtraceback()
             else:
                 value = local_value if local_value is not None else value
         return value
 
     def register_callback(self, callback, remove=False):
         """(Un)Register a callback
 
         Parameters
         ----------
         callback: method handle
             Method to be registered or unregistered.
         remove=False: bool
             Whether to unregister the callback."""
         
         # (Un)Register the callback.
         if remove and callback in self.callbacks:
             self.callbacks.remove(callback)
         elif not remove and callback not in self.callbacks:
             self.callbacks.append(callback)
 
 def _show_traceback(method):
     """decorator for showing tracebacks in IPython"""
     def m(self, *args, **kwargs):
         try:
             return(method(self, *args, **kwargs))
         except Exception as e:
             ip = get_ipython()
             if ip is None:
                 self.log.warn("Exception in widget method %s: %s", method, e, exc_info=True)
             else:
                 ip.showtraceback()
     return m
 
 class Widget(LoggingConfigurable):
     #-------------------------------------------------------------------------
     # Class attributes
     #-------------------------------------------------------------------------
     _widget_construction_callback = None
     widgets = {}
 
     @staticmethod
     def on_widget_constructed(callback):
         """Registers a callback to be called when a widget is constructed.
 
         The callback must have the following signature:
         callback(widget)"""
         Widget._widget_construction_callback = callback
 
     @staticmethod
     def _call_widget_constructed(widget):
         """Static method, called when a widget is constructed."""
         if Widget._widget_construction_callback is not None and callable(Widget._widget_construction_callback):
             Widget._widget_construction_callback(widget)
 
     #-------------------------------------------------------------------------
     # Traits
     #-------------------------------------------------------------------------
     _model_name = Unicode('WidgetModel', help="""Name of the backbone model 
         registered in the front-end to create and sync this widget with.""")
     _view_name = Unicode(help="""Default view registered in the front-end
         to use to represent the widget.""", sync=True)
     _comm = Instance('IPython.kernel.comm.Comm')
     
     msg_throttle = Int(3, sync=True, help="""Maximum number of msgs the 
         front-end can send before receiving an idle msg from the back-end.""")
     
     keys = List()
     def _keys_default(self):
         return [name for name in self.traits(sync=True)]
     
     _property_lock = Tuple((None, None))
     
     _display_callbacks = Instance(CallbackDispatcher, ())
     _msg_callbacks = Instance(CallbackDispatcher, ())
     
     #-------------------------------------------------------------------------
     # (Con/de)structor
     #-------------------------------------------------------------------------
     def __init__(self, **kwargs):
         """Public constructor"""
         super(Widget, self).__init__(**kwargs)
 
         self.on_trait_change(self._handle_property_changed, self.keys)
         Widget._call_widget_constructed(self)
 
     def __del__(self):
         """Object disposal"""
         self.close()
 
     #-------------------------------------------------------------------------
     # Properties
     #-------------------------------------------------------------------------
 
     @property
     def comm(self):
         """Gets the Comm associated with this widget.
 
         If a Comm doesn't exist yet, a Comm will be created automagically."""
         if self._comm is None:
             # Create a comm.
             self._comm = Comm(target_name=self._model_name)
             self._comm.on_msg(self._handle_msg)
-            self._comm.on_close(self._close)
             Widget.widgets[self.model_id] = self
 
             # first update
             self.send_state()
         return self._comm
     
     @property
     def model_id(self):
         """Gets the model id of this widget.
 
         If a Comm doesn't exist yet, a Comm will be created automagically."""
         return self.comm.comm_id
 
     #-------------------------------------------------------------------------
     # Methods
     #-------------------------------------------------------------------------
-    def _close(self):
-        """Private close - cleanup objects, registry entries"""
-        del Widget.widgets[self.model_id]
-        self._comm = None
 
     def close(self):
         """Close method.
 
-        Closes the widget which closes the underlying comm.
+        Closes the underlying comm.
         When the comm is closed, all of the widget views are automatically
         removed from the front-end."""
+        del Widget.widgets[self.model_id]
         if self._comm is not None:
             self._comm.close()
-            self._close()
-
+        self._comm = None
+    
     def send_state(self, key=None):
         """Sends the widget state, or a piece of it, to the front-end.
 
         Parameters
         ----------
         key : unicode (optional)
             A single property's name to sync with the front-end.
         """
         self._send({
             "method" : "update",
             "state"  : self.get_state()
         })
 
     def get_state(self, key=None):
         """Gets the widget state, or a piece of it.
 
         Parameters
         ----------
         key : unicode (optional)
             A single property's name to get.
         """
         keys = self.keys if key is None else [key]
         state = {}
         for k in keys:
             f = self.trait_metadata(k, 'to_json')
             if f is None:
                 f = self._trait_to_json
             value = getattr(self, k)
             state[k] = f(value)
         return state
     
     def send(self, content):
         """Sends a custom msg to the widget model in the front-end.
 
         Parameters
         ----------
         content : dict
             Content of the message to send.
         """
         self._send({"method": "custom", "content": content})
 
     def on_msg(self, callback, remove=False):
         """(Un)Register a custom msg receive callback.
 
         Parameters
         ----------
         callback: callable
             callback will be passed two arguments when a message arrives::
             
                 callback(widget, content)
             
         remove: bool
             True if the callback should be unregistered."""
         self._msg_callbacks.register_callback(callback, remove=remove)
 
     def on_displayed(self, callback, remove=False):
         """(Un)Register a widget displayed callback.
 
         Parameters
         ----------
         callback: method handler
             Must have a signature of::
             
                 callback(widget, **kwargs)
             
             kwargs from display are passed through without modification.
         remove: bool
             True if the callback should be unregistered."""
         self._display_callbacks.register_callback(callback, remove=remove)
 
     #-------------------------------------------------------------------------
     # Support methods
     #-------------------------------------------------------------------------
     @contextmanager
     def _lock_property(self, key, value):
         """Lock a property-value pair.
 
         NOTE: This, in addition to the single lock for all state changes, is 
         flawed.  In the future we may want to look into buffering state changes 
         back to the front-end."""
         self._property_lock = (key, value)
         try:
             yield
         finally:
             self._property_lock = (None, None)
 
     def _should_send_property(self, key, value):
         """Check the property lock (property_lock)"""
         return key != self._property_lock[0] or \
         value != self._property_lock[1]
     
     # Event handlers
     @_show_traceback
     def _handle_msg(self, msg):
         """Called when a msg is received from the front-end"""
         data = msg['content']['data']
         method = data['method']
         if not method in ['backbone', 'custom']:
             self.log.error('Unknown front-end to back-end widget msg with method "%s"' % method)
 
         # Handle backbone sync methods CREATE, PATCH, and UPDATE all in one.
         if method == 'backbone' and 'sync_data' in data:
             sync_data = data['sync_data']
             self._handle_receive_state(sync_data) # handles all methods
 
         # Handle a custom msg from the front-end
         elif method == 'custom':
             if 'content' in data:
                 self._handle_custom_msg(data['content'])
 
     def _handle_receive_state(self, sync_data):
         """Called when a state is received from the front-end."""
         for name in self.keys:
             if name in sync_data:
                 f = self.trait_metadata(name, 'from_json')
                 if f is None:
                     f = self._trait_from_json
                 value = f(sync_data[name])
                 with self._lock_property(name, value):
                     setattr(self, name, value)
 
     def _handle_custom_msg(self, content):
         """Called when a custom msg is received."""
         self._msg_callbacks(self, content)
 
     def _handle_property_changed(self, name, old, new):
         """Called when a property has been changed."""
         # Make sure this isn't information that the front-end just sent us.
         if self._should_send_property(name, new):
             # Send new state to front-end
             self.send_state(key=name)
 
     def _handle_displayed(self, **kwargs):
         """Called when a view has been displayed for this widget instance"""
         self._display_callbacks(self, **kwargs)
 
     def _trait_to_json(self, x):
         """Convert a trait value to json
 
         Traverse lists/tuples and dicts and serialize their values as well.
         Replace any widgets with their model_id
         """
         if isinstance(x, dict):
             return {k: self._trait_to_json(v) for k, v in x.items()}
         elif isinstance(x, (list, tuple)):
             return [self._trait_to_json(v) for v in x]
         elif isinstance(x, Widget):
             return "IPY_MODEL_" + x.model_id
         else:
             return x # Value must be JSON-able
 
     def _trait_from_json(self, x):
         """Convert json values to objects
 
         Replace any strings representing valid model id values to Widget references.
         """
         if isinstance(x, dict):
             return {k: self._trait_from_json(v) for k, v in x.items()}
         elif isinstance(x, (list, tuple)):
             return [self._trait_from_json(v) for v in x]
         elif isinstance(x, string_types) and x.startswith('IPY_MODEL_') and x[10:] in Widget.widgets:
             # we want to support having child widgets at any level in a hierarchy
             # trusting that a widget UUID will not appear out in the wild
             return Widget.widgets[x]
         else:
             return x
 
     def _ipython_display_(self, **kwargs):
         """Called when `IPython.display.display` is called on the widget."""
         # Show view.  By sending a display message, the comm is opened and the
         # initial state is sent.
         self._send({"method": "display"})
         self._handle_displayed(**kwargs)
 
     def _send(self, msg):
         """Sends a message to the model in the front-end."""
         self.comm.send(msg)
 
 
 class DOMWidget(Widget):
     visible = Bool(True, help="Whether the widget is visible.", sync=True)
     _css = List(sync=True) # Internal CSS property list: (selector, key, value)
 
     def get_css(self, key, selector=""):
         """Get a CSS property of the widget.
 
         Note: This function does not actually request the CSS from the 
         front-end;  Only properties that have been set with set_css can be read.
 
         Parameters
         ----------
         key: unicode
             CSS key
         selector: unicode (optional)
             JQuery selector used when the CSS key/value was set.
         """
         if selector in self._css and key in self._css[selector]:
             return self._css[selector][key]
         else:
             return None
 
     def set_css(self, dict_or_key, value=None, selector=''):
         """Set one or more CSS properties of the widget.
 
         This function has two signatures:
         - set_css(css_dict, selector='')
         - set_css(key, value, selector='')
 
         Parameters
         ----------
         css_dict : dict
             CSS key/value pairs to apply
         key: unicode
             CSS key
         value:
             CSS value
         selector: unicode (optional, kwarg only)
             JQuery selector to use to apply the CSS key/value.  If no selector 
             is provided, an empty selector is used.  An empty selector makes the 
             front-end try to apply the css to a default element.  The default
             element is an attribute unique to each view, which is a DOM element
             of the view that should be styled with common CSS (see 
             `$el_to_style` in the Javascript code).
         """
         if value is None:
             css_dict = dict_or_key
         else:
             css_dict = {dict_or_key: value}
         
         for (key, value) in css_dict.items():
             # First remove the selector/key pair from the css list if it exists.
             # Then add the selector/key pair and new value to the bottom of the 
             # list.
             self._css = [x for x in self._css if not (x[0]==selector and x[1]==key)]
             self._css += [(selector, key, value)]
         self.send_state('_css')
 
     def add_class(self, class_names, selector=""):
         """Add class[es] to a DOM element.
 
         Parameters
         ----------
         class_names: unicode or list
             Class name(s) to add to the DOM element(s).
         selector: unicode (optional)
             JQuery selector to select the DOM element(s) that the class(es) will
             be added to.
         """
         class_list = class_names
         if isinstance(class_list, (list, tuple)):
             class_list = ' '.join(class_list)
 
         self.send({
             "msg_type"   : "add_class",
             "class_list" : class_list,
             "selector"   : selector
         })
 
     def remove_class(self, class_names, selector=""):
         """Remove class[es] from a DOM element.
 
         Parameters
         ----------
         class_names: unicode or list
             Class name(s) to remove from  the DOM element(s).
         selector: unicode (optional)
             JQuery selector to select the DOM element(s) that the class(es) will
             be removed from.
         """
         class_list = class_names
         if isinstance(class_list, (list, tuple)):
             class_list = ' '.join(class_list)
 
         self.send({
             "msg_type"   : "remove_class",
             "class_list" : class_list,
             "selector"   : selector,
         })

IPython/kernel/comm/comm.py

@@ -1,133 +1,135 b''
 """Base class for a Comm"""
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 import uuid
 
 from IPython.config import LoggingConfigurable
 from IPython.core.getipython import get_ipython
 
 from IPython.utils.jsonutil import json_clean
 from IPython.utils.traitlets import Instance, Unicode, Bytes, Bool, Dict, Any
 
 
 class Comm(LoggingConfigurable):
     
     shell = Instance('IPython.core.interactiveshell.InteractiveShellABC')
     def _shell_default(self):
         return get_ipython()
     
     iopub_socket = Any()
     def _iopub_socket_default(self):
         return self.shell.kernel.iopub_socket
     session = Instance('IPython.kernel.zmq.session.Session')
     def _session_default(self):
         if self.shell is None:
             return
         return self.shell.kernel.session
     
     target_name = Unicode('comm')
     
     topic = Bytes()
     def _topic_default(self):
         return ('comm-%s' % self.comm_id).encode('ascii')
     
     _open_data = Dict(help="data dict, if any, to be included in comm_open")
     _close_data = Dict(help="data dict, if any, to be included in comm_close")
     
     _msg_callback = Any()
     _close_callback = Any()
     
     _closed = Bool(False)
     comm_id = Unicode()
     def _comm_id_default(self):
         return uuid.uuid4().hex
     
     primary = Bool(True, help="Am I the primary or secondary Comm?")
     
     def __init__(self, target_name='', data=None, **kwargs):
         if target_name:
             kwargs['target_name'] = target_name
         super(Comm, self).__init__(**kwargs)
-        get_ipython().comm_manager.register_comm(self)
         if self.primary:
             # I am primary, open my peer.
             self.open(data)
     
     def _publish_msg(self, msg_type, data=None, metadata=None, **keys):
         """Helper for sending a comm message on IOPub"""
         data = {} if data is None else data
         metadata = {} if metadata is None else metadata
         content = json_clean(dict(data=data, comm_id=self.comm_id, **keys))
         self.session.send(self.iopub_socket, msg_type,
             content,
             metadata=json_clean(metadata),
             parent=self.shell.get_parent(),
             ident=self.topic,
         )
     
     def __del__(self):
         """trigger close on gc"""
         self.close()
     
     # publishing messages
     
     def open(self, data=None, metadata=None):
         """Open the frontend-side version of this comm"""
         if data is None:
             data = self._open_data
+        self._closed = False
+        get_ipython().comm_manager.register_comm(self)
         self._publish_msg('comm_open', data, metadata, target_name=self.target_name)
     
     def close(self, data=None, metadata=None):
         """Close the frontend-side version of this comm"""
         if self._closed:
             # only close once
             return
         if data is None:
             data = self._close_data
         self._publish_msg('comm_close', data, metadata)
+        get_ipython().comm_manager.unregister_comm(self)
         self._closed = True
     
     def send(self, data=None, metadata=None):
         """Send a message to the frontend-side version of this comm"""
         self._publish_msg('comm_msg', data, metadata)
     
     # registering callbacks
     
     def on_close(self, callback):
         """Register a callback for comm_close
         
         Will be called with the `data` of the close message.
         
         Call `on_close(None)` to disable an existing callback.
         """
         self._close_callback = callback
     
     def on_msg(self, callback):
         """Register a callback for comm_msg
         
         Will be called with the `data` of any comm_msg messages.
         
         Call `on_msg(None)` to disable an existing callback.
         """
         self._msg_callback = callback
     
     # handling of incoming messages
     
     def handle_close(self, msg):
         """Handle a comm_close message"""
         self.log.debug("handle_close[%s](%s)", self.comm_id, msg)
         if self._close_callback:
             self._close_callback(msg)
     
     def handle_msg(self, msg):
         """Handle a comm_msg message"""
         self.log.debug("handle_msg[%s](%s)", self.comm_id, msg)
         if self._msg_callback:
             self.shell.events.trigger('pre_execute')
             self._msg_callback(msg)
             self.shell.events.trigger('post_execute')
 
 
 __all__ = ['Comm']

IPython/kernel/comm/manager.py

@@ -1,151 +1,148 b''
 """Base class to manage comms"""
 
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
 import sys
 
 from IPython.config import LoggingConfigurable
 from IPython.core.prompts import LazyEvaluate
 from IPython.core.getipython import get_ipython
 
 from IPython.utils.importstring import import_item
 from IPython.utils.py3compat import string_types
 from IPython.utils.traitlets import Instance, Unicode, Dict, Any
 
 from .comm import Comm
 
 
 def lazy_keys(dikt):
     """Return lazy-evaluated string representation of a dictionary's keys
     
     Key list is only constructed if it will actually be used.
     Used for debug-logging.
     """
     return LazyEvaluate(lambda d: list(d.keys()))
 
 
 class CommManager(LoggingConfigurable):
     """Manager for Comms in the Kernel"""
     
     shell = Instance('IPython.core.interactiveshell.InteractiveShellABC')
     def _shell_default(self):
         return get_ipython()
     iopub_socket = Any()
     def _iopub_socket_default(self):
         return self.shell.kernel.iopub_socket
     session = Instance('IPython.kernel.zmq.session.Session')
     def _session_default(self):
         if self.shell is None:
             return
         return self.shell.kernel.session
     
     comms = Dict()
     targets = Dict()
     
     # Public APIs
     
     def register_target(self, target_name, f):
         """Register a callable f for a given target name
         
         f will be called with two arguments when a comm_open message is received with `target`:
         
         - the Comm instance
         - the `comm_open` message itself.
         
         f can be a Python callable or an import string for one.
         """
         if isinstance(f, string_types):
             f = import_item(f)
         
         self.targets[target_name] = f
     
     def unregister_target(self, target_name, f):
         """Unregister a callable registered with register_target"""
         return self.targets.pop(target_name);
     
     def register_comm(self, comm):
         """Register a new comm"""
         comm_id = comm.comm_id
         comm.shell = self.shell
         comm.iopub_socket = self.iopub_socket
         self.comms[comm_id] = comm
         return comm_id
     
-    def unregister_comm(self, comm_id):
+    def unregister_comm(self, comm):
         """Unregister a comm, and close its counterpart"""
         # unlike get_comm, this should raise a KeyError
-        comm = self.comms.pop(comm_id)
-        comm.close()
+        comm = self.comms.pop(comm.comm_id)
     
     def get_comm(self, comm_id):
         """Get a comm with a particular id
         
         Returns the comm if found, otherwise None.
         
         This will not raise an error,
         it will log messages if the comm cannot be found.
         """
         if comm_id not in self.comms:
             self.log.error("No such comm: %s", comm_id)
             self.log.debug("Current comms: %s", lazy_keys(self.comms))
             return
         # call, because we store weakrefs
         comm = self.comms[comm_id]
         return comm
     
     # Message handlers
     def comm_open(self, stream, ident, msg):
         """Handler for comm_open messages"""
         content = msg['content']
         comm_id = content['comm_id']
         target_name = content['target_name']
         f = self.targets.get(target_name, None)
         comm = Comm(comm_id=comm_id,
                     shell=self.shell,
                     iopub_socket=self.iopub_socket,
                     primary=False,
         )
         if f is None:
             self.log.error("No such comm target registered: %s", target_name)
             comm.close()
             return
-        self.register_comm(comm)
         try:
             f(comm, msg)
         except Exception:
             self.log.error("Exception opening comm with target: %s", target_name, exc_info=True)
             comm.close()
-            self.unregister_comm(comm_id)
     
     def comm_msg(self, stream, ident, msg):
         """Handler for comm_msg messages"""
         content = msg['content']
         comm_id = content['comm_id']
         comm = self.get_comm(comm_id)
         if comm is None:
             # no such comm
             return
         try:
             comm.handle_msg(msg)
         except Exception:
             self.log.error("Exception in comm_msg for %s", comm_id, exc_info=True)
     
     def comm_close(self, stream, ident, msg):
         """Handler for comm_close messages"""
         content = msg['content']
         comm_id = content['comm_id']
         comm = self.get_comm(comm_id)
         if comm is None:
             # no such comm
             self.log.debug("No such comm to close: %s", comm_id)
             return
         del self.comms[comm_id]
         
         try:
             comm.handle_close(msg)
         except Exception:
             self.log.error("Exception handling comm_close for %s", comm_id, exc_info=True)
 
 
 __all__ = ['CommManager']