unregister and register both take comm as argument

IPython/kernel/comm/manager.py

@@ -72,10 +72,10 b' class CommManager(LoggingConfigurable):'
         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 = self.comms.pop(comm.comm_id)
     
     def get_comm(self, comm_id):
         """Get a comm with a particular id
@@ -115,7 +115,7 b' class CommManager(LoggingConfigurable):'
         except Exception:
             self.log.error("Exception opening comm with target: %s", target_name, exc_info=True)
             comm.close()
-            self.unregister_comm(comm_id)
+            self.unregister_comm(comm)
     
     def comm_msg(self, stream, ident, msg):
         """Handler for comm_msg messages"""