comm.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  1. """Base class for a Comm"""
  2. # Copyright (c) IPython Development Team.
  3. # Distributed under the terms of the Modified BSD License.
  4. import uuid
  5. from traitlets.config import LoggingConfigurable
  6. from ipykernel.kernelbase import Kernel
  7. from ipykernel.jsonutil import json_clean
  8. from traitlets import Instance, Unicode, Bytes, Bool, Dict, Any, default
  9. class Comm(LoggingConfigurable):
  10. """Class for communicating between a Frontend and a Kernel"""
  11. kernel = Instance('ipykernel.kernelbase.Kernel', allow_none=True)
  12. @default('kernel')
  13. def _default_kernel(self):
  14. if Kernel.initialized():
  15. return Kernel.instance()
  16. comm_id = Unicode()
  17. @default('comm_id')
  18. def _default_comm_id(self):
  19. return uuid.uuid4().hex
  20. primary = Bool(True, help="Am I the primary or secondary Comm?")
  21. target_name = Unicode('comm')
  22. target_module = Unicode(None, allow_none=True, help="""requirejs module from
  23. which to load comm target.""")
  24. topic = Bytes()
  25. @default('topic')
  26. def _default_topic(self):
  27. return ('comm-%s' % self.comm_id).encode('ascii')
  28. _open_data = Dict(help="data dict, if any, to be included in comm_open")
  29. _close_data = Dict(help="data dict, if any, to be included in comm_close")
  30. _msg_callback = Any()
  31. _close_callback = Any()
  32. _closed = Bool(True)
  33. def __init__(self, target_name='', data=None, metadata=None, buffers=None, **kwargs):
  34. if target_name:
  35. kwargs['target_name'] = target_name
  36. super(Comm, self).__init__(**kwargs)
  37. if self.kernel:
  38. if self.primary:
  39. # I am primary, open my peer.
  40. self.open(data=data, metadata=metadata, buffers=buffers)
  41. else:
  42. self._closed = False
  43. def _publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
  44. """Helper for sending a comm message on IOPub"""
  45. data = {} if data is None else data
  46. metadata = {} if metadata is None else metadata
  47. content = json_clean(dict(data=data, comm_id=self.comm_id, **keys))
  48. self.kernel.session.send(self.kernel.iopub_socket, msg_type,
  49. content,
  50. metadata=json_clean(metadata),
  51. parent=self.kernel._parent_header,
  52. ident=self.topic,
  53. buffers=buffers,
  54. )
  55. def __del__(self):
  56. """trigger close on gc"""
  57. self.close()
  58. # publishing messages
  59. def open(self, data=None, metadata=None, buffers=None):
  60. """Open the frontend-side version of this comm"""
  61. if data is None:
  62. data = self._open_data
  63. comm_manager = getattr(self.kernel, 'comm_manager', None)
  64. if comm_manager is None:
  65. raise RuntimeError("Comms cannot be opened without a kernel "
  66. "and a comm_manager attached to that kernel.")
  67. comm_manager.register_comm(self)
  68. try:
  69. self._publish_msg('comm_open',
  70. data=data, metadata=metadata, buffers=buffers,
  71. target_name=self.target_name,
  72. target_module=self.target_module,
  73. )
  74. self._closed = False
  75. except:
  76. comm_manager.unregister_comm(self)
  77. raise
  78. def close(self, data=None, metadata=None, buffers=None):
  79. """Close the frontend-side version of this comm"""
  80. if self._closed:
  81. # only close once
  82. return
  83. self._closed = True
  84. # nothing to send if we have no kernel
  85. # can be None during interpreter cleanup
  86. if not self.kernel:
  87. return
  88. if data is None:
  89. data = self._close_data
  90. self._publish_msg('comm_close',
  91. data=data, metadata=metadata, buffers=buffers,
  92. )
  93. self.kernel.comm_manager.unregister_comm(self)
  94. def send(self, data=None, metadata=None, buffers=None):
  95. """Send a message to the frontend-side version of this comm"""
  96. self._publish_msg('comm_msg',
  97. data=data, metadata=metadata, buffers=buffers,
  98. )
  99. # registering callbacks
  100. def on_close(self, callback):
  101. """Register a callback for comm_close
  102. Will be called with the `data` of the close message.
  103. Call `on_close(None)` to disable an existing callback.
  104. """
  105. self._close_callback = callback
  106. def on_msg(self, callback):
  107. """Register a callback for comm_msg
  108. Will be called with the `data` of any comm_msg messages.
  109. Call `on_msg(None)` to disable an existing callback.
  110. """
  111. self._msg_callback = callback
  112. # handling of incoming messages
  113. def handle_close(self, msg):
  114. """Handle a comm_close message"""
  115. self.log.debug("handle_close[%s](%s)", self.comm_id, msg)
  116. if self._close_callback:
  117. self._close_callback(msg)
  118. def handle_msg(self, msg):
  119. """Handle a comm_msg message"""
  120. self.log.debug("handle_msg[%s](%s)", self.comm_id, msg)
  121. if self._msg_callback:
  122. shell = self.kernel.shell
  123. if shell:
  124. shell.events.trigger('pre_execute')
  125. self._msg_callback(msg)
  126. if shell:
  127. shell.events.trigger('post_execute')
  128. __all__ = ['Comm']