Last active
July 2, 2024 06:00
-
-
Save KaoruNishikawa/9ecb902e6b280df09273d5a88c22d3b6 to your computer and use it in GitHub Desktop.
Run ipynb cells one by one, extracting variable snapshot at each steps
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| """ | |
| ## Limitations | |
| - Extracted variables are just string representation, not arbitrary Python objects | |
| - Variables defined by substitution (e.g. `a = 1`) would be extracted, but others (e.g. `def func():`) won't | |
| """ | |
| import logging | |
| import time | |
| import nbformat | |
| from jupyter_client.channels import AsyncZMQSocketChannel | |
| from jupyter_client.session import Session | |
| from nbclient import NotebookClient | |
| logger = logging.getLogger(__name__) | |
| seq = 0 | |
| nb = nbformat.read('path/to/my.ipynb', as_version=4) | |
| client = NotebookClient( | |
| nb, | |
| timeout=5, | |
| kernel_name="python3", | |
| allow_errors=True, | |
| resources={"metadata": {"path": "./"}}, | |
| ) | |
| def get_variables( | |
| session: Session, | |
| socket: AsyncZMQSocketChannel, | |
| timeout: int | float = 5, | |
| ) -> dict[str, dict]: | |
| global seq | |
| seq += 1 | |
| content = {"type": "request", "command": "inspectVariables", "seq": seq} | |
| msg = session.msg("debug_request", content) | |
| msg_id = msg["header"]["msg_id"] | |
| session.send(socket, msg) | |
| start = time.monotonic() | |
| while time.monotonic() - start < timeout: | |
| idents, res = session.recv(socket) | |
| if (idents is not None) and (res["parent_header"]["msg_id"] == msg_id): | |
| logger.debug(res) | |
| return res["content"]["body"]["variables"] | |
| time.sleep(0.1) | |
| raise TimeoutError("Timeout asking for notebook variables") | |
| with client.setup_kernel(): | |
| for i, cell in enumerate(nb.cells): | |
| _result = client.execute_cell(cell, i) # Cell execution result, including its output | |
| reply = get_variables( | |
| client.kc.session, client.kc.control_channel.socket, timeout=client.timeout | |
| ) | |
| for v in reply: | |
| print(v) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment