Skip to content

Commit 3c7b60b

Browse files
committed
DOC: Add docstrings to CellWatcher class
1 parent cf6fe9d commit 3c7b60b

File tree

1 file changed

+131
-39
lines changed

1 file changed

+131
-39
lines changed

itkwidgets/cell_watcher.py

Lines changed: 131 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
import asyncio
22
import sys
33
from inspect import isawaitable, iscoroutinefunction
4-
from typing import Dict, List
4+
from typing import Callable, Dict, List
55
from IPython import get_ipython
6+
from IPython.core.interactiveshell import ExecutionResult
67
from queue import Queue
78
from imjoy_rpc.utils import FuturePromise
9+
from zmq.eventloop.zmqstream import ZMQStream
810

911
background_tasks = set()
1012

@@ -72,13 +74,23 @@ def viewer_ready(self, view: str) -> bool:
7274

7375

7476
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+
7582
def __new__(cls):
76-
if not hasattr(cls, '_instance'):
83+
"""Create a singleton class."""
84+
if not hasattr(cls, "_instance"):
7785
cls._instance = super(CellWatcher, cls).__new__(cls)
7886
cls._instance.setup()
7987
return cls._instance
8088

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+
"""
8294
self.viewers = Viewers()
8395
self.shell = get_ipython()
8496
self.kernel = self.shell.kernel
@@ -102,22 +114,52 @@ def setup(self):
102114

103115
# Call self.post_run_cell every time the post_run_cell signal is emitted
104116
# 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.
106121
107-
def add_viewer(self, view):
122+
:param view: The unique string identifier for the Viewer object
123+
:type view: str
124+
"""
108125
# Track all Viewer instances
109126
self.viewers.add_viewer(view)
110127

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+
"""
112139
self.viewers.update_viewer_status(view, status)
113-
if self.waiting_on_viewer:
140+
if status and self.waiting_on_viewer:
114141
# Might be ready now, try again
115142
self.create_task(self.execute_next_request)
116143

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+
"""
118155
return self.viewers.viewer_ready(view)
119156

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+
"""
121163
global background_tasks
122164
try:
123165
# "Handle" exceptions here to prevent further errors. Exceptions
@@ -127,40 +169,83 @@ def _task_cleanup(self, task):
127169
except:
128170
background_tasks.discard(task)
129171

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+
"""
131178
global background_tasks
132179
# The event loop only keeps weak references to tasks.
133180
# Gather them into a set to avoid garbage collection mid-task.
134181
task = asyncio.create_task(fn())
135182
background_tasks.add(task)
136183
task.add_done_callback(self._task_cleanup)
137184

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+
"""
139198
self._events.put((stream, ident, parent))
140199
if self._events.qsize() == 1 and self.ready_to_run_next_cell():
141200
# We've added a new task to an empty queue.
142201
# Begin executing tasks again.
143202
self.create_task(self.execute_next_request)
144203

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+
"""
146220
# ipykernel 6+
147221
self.capture_event(stream, ident, parent)
148222

149223
@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+
"""
152231
getters_resolved = [f.done() for f in self.results.values()]
153232
return all(getters_resolved)
154233

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+
"""
158241
self.waiting_on_viewer = len(self.viewers.not_created)
159242
return self.all_getters_resolved and not self.waiting_on_viewer
160243

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+
"""
164249
if self._events.empty():
165250
self.abort_all = False
166251

@@ -172,7 +257,8 @@ async def execute_next_request(self):
172257
# Continue processing the remaining queued tasks
173258
await self._execute_next_request()
174259

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"""
176262
# Here we actually run the queued cell as it would have been run
177263
stream, ident, parent = self.current_request
178264

@@ -202,32 +288,38 @@ async def _execute_next_request(self):
202288
# Continue processing the remaining queued tasks
203289
self.create_task(self.execute_next_request)
204290

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"""
207293
# FIXME: This is a temporary "fix" and does not handle updating output
208294
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+
"""
224310
# After each getter/setter resolves check if they've all resolved
225311
if self.all_getters_resolved:
226312
self.update_namespace()
227313
self.current_request = None
228314
self.create_task(self.execute_next_request)
229315

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+
"""
231323
# Abort remaining cells on error in execution
232324
if response.error_in_exec is not None:
233325
self.abort_all = True

0 commit comments

Comments
 (0)