-
Notifications
You must be signed in to change notification settings - Fork 0
/
jupyter_pstate.py
125 lines (111 loc) · 3.54 KB
/
jupyter_pstate.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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
import os, sys, types, inspect
from IPython import get_ipython
from IPython.core.interactiveshell import InteractiveShell
import nbformat
import pickle
import warnings
class ImportBreak(ImportError):
"""Exception for breaking import"""
pass
class NotebookFinder(object):
"""Module finder that locates IPython Notebooks"""
def __init__(self):
self.loaders = {}
@classmethod
def find_notebook(cls, fullname, path=None):
"""find a notebook, given its fully qualified name and an optional path"""
name = fullname.rsplit('.', 1)[-1]
if not path:
path = ['']
for d in path:
nb_path = os.path.join(d, name + ".ipynb")
if os.path.isfile(nb_path):
return nb_path
elif os.path.isfile(nb_path.replace("_", " ")):
return nb_path.replace("_", " ")
elif os.path.isfile(nb_path.replace("_", "-")):
return nb_path.replace("_", "-")
def find_module(self, fullname, path=None):
nb_path = self.find_notebook(fullname, path)
if not nb_path:
return
key = path
if path:
key = os.path.sep.join(path)
if key not in self.loaders:
self.loaders[key] = NotebookLoader(path)
return self.loaders[key]
class NotebookLoader(object):
"""Module Loader for IPython Notebooks"""
find_notebook = NotebookFinder.find_notebook
def __init__(self, path=None):
self.shell = InteractiveShell.instance()
self.path = path
def exceptImportBreak(self, mod, cells): # module name is used in pstate.ImportBreak
"""Transform and exec module code"""
try:
for cell in cells:
if cell.cell_type == 'code':
code = self.shell.input_transformer_manager.transform_cell(cell.source)
exec(code, mod.__dict__)
except ImportBreak:
return mod
return mod
def load_module(self, fullname):
"""import a notebook as a module"""
path = self.find_notebook(fullname, self.path)
nb = nbformat.read(path, as_version=4)
mod = types.ModuleType(fullname)
mod.__file__ = path
mod.__loader__ = self
mod.__dict__['get_ipython'] = get_ipython
sys.modules[fullname] = mod
save_user_ns = self.shell.user_ns
self.shell.user_ns = mod.__dict__
try:
mod = self.exceptImportBreak(mod, nb.cells)
finally:
self.shell.user_ns = save_user_ns
return mod
class pstate(object):
"""Saving and loading program state with pickle.
STATE = pstate('./file_name.tmp')"""
def __init__(self, file_name):
self.file_name = file_name
def __call__(self, file_name):
return pstate(file_name)
def save(self, *args, pause=True, import_break=True):
"""Saving state.
Args:
*args - vars for pickle.dump
pause=True - stop code execution
import_break=True - stop code execution when importing
Examples:
file_name = STATE.save(val1, val2)
STATE.save(locals())"""
if import_break:
self.ImportBreak()
with open(self.file_name, 'wb') as file:
pickle.dump(args, file)
if pause:
print('State saved in:\n{}'.format(self.file_name))
input("Press Enter to continue...")
return self.file_name
def load(self):
"""Loading saved state.
val1, = STATE.load()
locals().update(STATE.load()[0])"""
with open(self.file_name, 'rb') as file:
return pickle.load(file)
def ImportBreak(self):
"""Break when importing as module."""
for stack in inspect.stack():
if stack[0].f_code.co_name == 'exceptImportBreak':
warnings.warn('Breaking import in this module!', UserWarning, stacklevel=2)
raise ImportBreak()
return None
if __name__ == '__main__':
pass
else:
# Autoadd NotebookFinder when importing.
sys.meta_path.append(NotebookFinder())