1
1
import asyncio
2
2
import sys
3
3
from inspect import isawaitable , iscoroutinefunction
4
- from typing import Dict , List
4
+ from typing import Callable , Dict , List
5
5
from IPython import get_ipython
6
+ from IPython .core .interactiveshell import ExecutionResult
6
7
from queue import Queue
7
8
from imjoy_rpc .utils import FuturePromise
9
+ from zmq .eventloop .zmqstream import ZMQStream
8
10
9
11
background_tasks = set ()
10
12
@@ -72,13 +74,23 @@ def viewer_ready(self, view: str) -> bool:
72
74
73
75
74
76
class CellWatcher (object ):
77
+ """A singleton class used in interactive Jupyter notebooks in order to
78
+ support asynchronous network communication that would otherwise be blocked
79
+ by the IPython kernel.
80
+ """
81
+
75
82
def __new__ (cls ):
76
- if not hasattr (cls , '_instance' ):
83
+ """Create a singleton class."""
84
+ if not hasattr (cls , "_instance" ):
77
85
cls ._instance = super (CellWatcher , cls ).__new__ (cls )
78
86
cls ._instance .setup ()
79
87
return cls ._instance
80
88
81
- def setup (self ):
89
+ def setup (self ) -> None :
90
+ """Perform the initial setup, including intercepting 'execute_request'
91
+ handlers so that we can handle them internally before the IPython
92
+ kernel does.
93
+ """
82
94
self .viewers = Viewers ()
83
95
self .shell = get_ipython ()
84
96
self .kernel = self .shell .kernel
@@ -102,22 +114,52 @@ def setup(self):
102
114
103
115
# Call self.post_run_cell every time the post_run_cell signal is emitted
104
116
# post_run_cell runs after interactive execution (e.g. a cell in a notebook)
105
- self .shell .events .register ('post_run_cell' , self .post_run_cell )
117
+ self .shell .events .register ("post_run_cell" , self .post_run_cell )
118
+
119
+ def add_viewer (self , view : str ) -> None :
120
+ """Add a new Viewer object to track.
106
121
107
- def add_viewer (self , view ):
122
+ :param view: The unique string identifier for the Viewer object
123
+ :type view: str
124
+ """
108
125
# Track all Viewer instances
109
126
self .viewers .add_viewer (view )
110
127
111
- def update_viewer_status (self , view , status ):
128
+ def update_viewer_status (self , view : str , status : bool ) -> None :
129
+ """Update a Viewer's 'ready' status. If the last cell run failed
130
+ because the viewer was unavailable try to run the cell again.
131
+
132
+ :param view: The unique string identifier for the Viewer object
133
+ :type view: str
134
+ :param status: Boolean value indicating whether or not the viewer is
135
+ available for requests or updates. This should be false when the plugin
136
+ API is not yet available or new data is not yet rendered.
137
+ :type status: bool
138
+ """
112
139
self .viewers .update_viewer_status (view , status )
113
- if self .waiting_on_viewer :
140
+ if status and self .waiting_on_viewer :
114
141
# Might be ready now, try again
115
142
self .create_task (self .execute_next_request )
116
143
117
- def viewer_ready (self , view ):
144
+ def viewer_ready (self , view : str ) -> bool :
145
+ """Request the 'ready' status of a viewer.
146
+
147
+ :param view: The unique string identifier for the Viewer object
148
+ :type view: str
149
+
150
+ :return: Boolean value indicating whether or not the viewer is
151
+ available for requests or updates. This will be false when the plugin
152
+ API is not yet available or new data is not yet rendered.
153
+ :rtype: bool
154
+ """
118
155
return self .viewers .viewer_ready (view )
119
156
120
- def _task_cleanup (self , task ):
157
+ def _task_cleanup (self , task : asyncio .Task ) -> None :
158
+ """Callback to discard references to tasks once they've completed.
159
+
160
+ :param task: Completed task that no longer needs a strong reference
161
+ :type task: asyncio.Task
162
+ """
121
163
global background_tasks
122
164
try :
123
165
# "Handle" exceptions here to prevent further errors. Exceptions
@@ -127,40 +169,83 @@ def _task_cleanup(self, task):
127
169
except :
128
170
background_tasks .discard (task )
129
171
130
- def create_task (self , fn ):
172
+ def create_task (self , fn : Callable ) -> None :
173
+ """Create a task from the function passed in.
174
+
175
+ :param fn: Coroutine to run concurrently as a Task
176
+ :type fn: Callable
177
+ """
131
178
global background_tasks
132
179
# The event loop only keeps weak references to tasks.
133
180
# Gather them into a set to avoid garbage collection mid-task.
134
181
task = asyncio .create_task (fn ())
135
182
background_tasks .add (task )
136
183
task .add_done_callback (self ._task_cleanup )
137
184
138
- def capture_event (self , stream , ident , parent ):
185
+ def capture_event (self , stream : ZMQStream , ident : list , parent : dict ) -> None :
186
+ """Capture execute_request messages so that we can queue and process
187
+ them concurrently as tasks to prevent blocking.
188
+
189
+ :param stream: Class to manage event-based messaging on a zmq socket
190
+ :type stream: ZMQStream
191
+ :param ident: ZeroMQ routing prefix, which can be zero or more socket
192
+ identities
193
+ :type ident: list
194
+ :param parent: A dictonary of dictionaries representing a complete
195
+ message as defined by the Jupyter message specification
196
+ :type parent: dict
197
+ """
139
198
self ._events .put ((stream , ident , parent ))
140
199
if self ._events .qsize () == 1 and self .ready_to_run_next_cell ():
141
200
# We've added a new task to an empty queue.
142
201
# Begin executing tasks again.
143
202
self .create_task (self .execute_next_request )
144
203
145
- async def capture_event_async (self , stream , ident , parent ):
204
+ async def capture_event_async (
205
+ self , stream : ZMQStream , ident : list , parent : dict
206
+ ) -> None :
207
+ """Capture execute_request messages so that we can queue and process
208
+ them concurrently as tasks to prevent blocking.
209
+ Asynchronous for ipykernel 6+.
210
+
211
+ :param stream: Class to manage event-based messaging on a zmq socket
212
+ :type stream: ZMQStream
213
+ :param ident: ZeroMQ routing prefix, which can be zero or more socket
214
+ identities
215
+ :type ident: list
216
+ :param parent: A dictonary of dictionaries representing a complete
217
+ message as defined by the Jupyter message specification
218
+ :type parent: dict
219
+ """
146
220
# ipykernel 6+
147
221
self .capture_event (stream , ident , parent )
148
222
149
223
@property
150
- def all_getters_resolved (self ):
151
- # Check if all of the getter/setter futures have resolved
224
+ def all_getters_resolved (self ) -> bool :
225
+ """Determine if all tasks representing asynchronous network calls that
226
+ fetch values have resolved.
227
+
228
+ :return: Whether or not all tasks for the current cell have resolved
229
+ :rtype: bool
230
+ """
152
231
getters_resolved = [f .done () for f in self .results .values ()]
153
232
return all (getters_resolved )
154
233
155
- def ready_to_run_next_cell (self ):
156
- # Any itk_viewer objects need to be available and all getters/setters
157
- # need to be resolved
234
+ def ready_to_run_next_cell (self ) -> bool :
235
+ """Determine if we are ready to run the next cell in the queue.
236
+
237
+ :return: If created Viewer objects are available and all futures are
238
+ resolved.
239
+ :rtype: bool
240
+ """
158
241
self .waiting_on_viewer = len (self .viewers .not_created )
159
242
return self .all_getters_resolved and not self .waiting_on_viewer
160
243
161
- async def execute_next_request (self ):
162
- # Modeled after the approach used in jupyter-ui-poll
163
- # https://github.com/Kirill888/jupyter-ui-poll/blob/f65b81f95623c699ed7fd66a92be6d40feb73cde/jupyter_ui_poll/_poll.py#L75-L101
244
+ async def execute_next_request (self ) -> None :
245
+ """Grab the next request if needed and then run the cell if it it ready
246
+ to be run. Modeled after the approach used in jupyter-ui-poll.
247
+ :ref: https://github.com/Kirill888/jupyter-ui-poll/blob/f65b81f95623c699ed7fd66a92be6d40feb73cde/jupyter_ui_poll/_poll.py#L75-L101
248
+ """
164
249
if self ._events .empty ():
165
250
self .abort_all = False
166
251
@@ -172,7 +257,8 @@ async def execute_next_request(self):
172
257
# Continue processing the remaining queued tasks
173
258
await self ._execute_next_request ()
174
259
175
- async def _execute_next_request (self ):
260
+ async def _execute_next_request (self ) -> None :
261
+ """Run the cell with the ipykernel shell_handler for execute_request"""
176
262
# Here we actually run the queued cell as it would have been run
177
263
stream , ident , parent = self .current_request
178
264
@@ -202,32 +288,38 @@ async def _execute_next_request(self):
202
288
# Continue processing the remaining queued tasks
203
289
self .create_task (self .execute_next_request )
204
290
205
- def update_namespace (self ):
206
- # Update the namespace variables with the results from the getters
291
+ def update_namespace (self ) -> None :
292
+ """ Update the namespace variables with the results from the getters"""
207
293
# FIXME: This is a temporary "fix" and does not handle updating output
208
294
keys = [k for k in self .shell .user_ns .keys ()]
209
- try :
210
- for key in keys :
211
- value = self . shell . user_ns [ key ]
212
- if asyncio . isfuture ( value ) and ( isinstance (value , FuturePromise ) or isinstance (value , asyncio .Task )):
213
- # Getters/setters return futures
214
- # They should all be resolved now, so use the result
215
- self . shell . user_ns [ key ] = value . result ()
216
- self . results . clear ()
217
- except Exception as e :
218
- self .results .clear ()
219
- self . abort_all = True
220
- self . create_task ( self . _execute_next_request )
221
- raise e
222
-
223
- def _callback ( self , * args , ** kwargs ):
295
+ for key in keys :
296
+ value = self . shell . user_ns [ key ]
297
+ if asyncio . isfuture ( value ) and (
298
+ isinstance (value , FuturePromise ) or isinstance (value , asyncio .Task )
299
+ ):
300
+ # Functions that need to return values from asynchronous
301
+ # network requests return futures. They should all be resolved
302
+ # now, so use the result.
303
+ self . shell . user_ns [ key ] = value . result ()
304
+ self .results .clear ()
305
+
306
+ def _callback ( self , * args , ** kwargs ) -> None :
307
+ """After each future resolves check to see if they are all resolved. If
308
+ so, update the namespace and run the next cell in the queue.
309
+ """
224
310
# After each getter/setter resolves check if they've all resolved
225
311
if self .all_getters_resolved :
226
312
self .update_namespace ()
227
313
self .current_request = None
228
314
self .create_task (self .execute_next_request )
229
315
230
- def post_run_cell (self , response ):
316
+ def post_run_cell (self , response : ExecutionResult ) -> None :
317
+ """Runs after interactive execution (e.g. a cell in a notebook). Set
318
+ the abort flag if there are errors produced by cell execution.
319
+
320
+ :param response: The response message produced by cell execution
321
+ :type response: ExecutionResult
322
+ """
231
323
# Abort remaining cells on error in execution
232
324
if response .error_in_exec is not None :
233
325
self .abort_all = True
0 commit comments