-
Notifications
You must be signed in to change notification settings - Fork 3
/
import_wrapper.py
105 lines (89 loc) · 3.7 KB
/
import_wrapper.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
from functools import partial
import re
import sys
try:
import colorama
colorama.init()
pre = (colorama.Fore.CYAN + colorama.Style.BRIGHT +
"AutoImport: " + colorama.Style.NORMAL +
colorama.Fore.WHITE)
except:
pre = "AutoImport: "
common_others = {"pd": "pandas", "np": "numpy", "sp": "scipy"}
class pipUnsuccessfulException(Exception):
pass
def custom_exc(ipython, shell, etype, evalue, tb, tb_offset=None):
shell.showtraceback((etype, evalue, tb), tb_offset)
while tb.tb_next:
tb = tb.tb_next
if not re.match("\\A<ipython-input-.*>\\Z", tb.tb_frame.f_code.co_filename):
# Innermost frame is not the IPython interactive environment.
return
try:
# Get the name of the module you tried to import
results = re.match("\\Aname '(.*)' is not defined\\Z", str(evalue))
if not results:
return
name = results.group(1)
custom_exc.last_name = name
try:
__import__(name)
except:
if common_others.get(name):
new_name = common_others.get(name)
try:
__import__(new_name)
r = ipython.ask_yes_no(
pre +
"{0} isn't a module, but {1} is. "
"Import {1} as {0}? (Y/n)".format(name, new_name))
if r:
name = "{} as {}".format(new_name, name)
else:
return
except Exception as e:
print(pre +
"{} isn't a module and nor is {}"
.format(name, new_name))
print(e)
return
else:
print(pre + "{} isn't a module".format(name))
try:
last_name = custom_exc.last_name
if ipython.ask_yes_no(pre + "Attempt to pip-install {}? (Y/n)"
.format(last_name)):
if __import__("pip").main(["install", last_name]) != 0:
raise pipUnsuccessfulException
else:
return
print(pre + "Installation completed successfully, importing...")
try:
res = ipython.run_cell("import {}".format(last_name))
print(pre + "Imported referenced module {}".format(last_name))
except:
print(pre + "{} isn't a module".format(last_name))
except pipUnsuccessfulException:
print(pre + "Installation with pip failed")
except AttributeError:
print(pre + "No module to install")
except ImportError:
print(pre + "pip not found")
return
# Import the module
ipython.run_code("import {}".format(name))
print(pre + "Imported referenced module {!r}, will retry".format(name))
print("".join("-" for _ in range(75)))
except Exception as e:
print(pre + ("Attempted to import {!r}"
"but an exception occured".format(name)))
try:
# Run the failed line again
res = ipython.run_cell(
list(ipython.history_manager.get_range())[-1][-1])
except Exception as e:
print(pre + "Another exception occured while retrying")
shell.showtraceback((type(e), e, None), None)
def load_ipython_extension(ipython):
# Bind the function we created to IPython's exception handler
ipython.set_custom_exc((NameError,), partial(custom_exc, ipython))