-
Notifications
You must be signed in to change notification settings - Fork 286
/
addRunCellHook.py
40 lines (32 loc) · 1.67 KB
/
addRunCellHook.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import types as _VSCODE_types
import os
import hashlib as _VSCODE_hashlib
from IPython import get_ipython
# Variable that prevent wrapping from happing more than_once
__VSCODE_wrapped_run_cell = False
# This function computes the hash for the code. It must follow the same algorithm as in use here:
# https://github.com/microsoft/vscode-jupyter/blob/312511f3cbd8b2bb5bc70fa9b771429e22d0c258/src/client/datascience/editor-integration/cellhashprovider.ts#L181
def __VSCODE_compute_hash(code, number=0):
hash_digest = _VSCODE_hashlib.sha1(code.encode("utf-8")).hexdigest()
return "<ipython-input-{0}-{1}>".format(number, hash_digest[:12])
# This function will wrap InteractiveShell.run-cell and force its IPYKERNEL_CELL_NAME to match the
# predicted value used by the Jupyter Extension
def __VSCODE_wrap_run_cell(wrapped_func):
old_func = wrapped_func.__func__
__VSCODE_wrapped_run_cell = True
def wrapper(*args, **kwargs):
try:
store_history = args[2] if len(args) > 2 else kwargs["store_history"]
if store_history:
predicted_name = __VSCODE_compute_hash(args[1], args[0].execution_count)
os.environ["IPYKERNEL_CELL_NAME"] = predicted_name
result = old_func(*args, **kwargs)
if store_history:
del os.environ["IPYKERNEL_CELL_NAME"]
return result
except:
return old_func(*args, **kwargs)
return _VSCODE_types.MethodType(wrapper, wrapped_func.__self__)
# Double check this still exists on IPython
if get_ipython() and get_ipython().run_cell and not __VSCODE_wrapped_run_cell:
get_ipython().run_cell = __VSCODE_wrap_run_cell(get_ipython().run_cell)