Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/IPython/extensions/storemagic.py: 16%

1""" 

2%store magic for lightweight persistence. 

3 

4Stores variables, aliases and macros in IPython's database. 

5 

6To automatically restore stored variables at startup, add this to your 

7:file:`ipython_config.py` file:: 

8 

9 c.StoreMagics.autorestore = True 

10""" 

11 

12# Copyright (c) IPython Development Team. 

13# Distributed under the terms of the Modified BSD License. 

14 

15import inspect, os, sys, textwrap 

16 

17from IPython.core.error import UsageError 

18from IPython.core.magic import Magics, magics_class, line_magic 

19from IPython.testing.skipdoctest import skip_doctest 

20from traitlets import Bool 

21 

22 

23def restore_aliases(ip, alias=None): 

24 staliases = ip.db.get('stored_aliases', {}) 

25 if alias is None: 

26 for k,v in staliases.items(): 

27 # print("restore alias",k,v) # dbg 

28 #self.alias_table[k] = v 

29 ip.alias_manager.define_alias(k,v) 

30 else: 

31 ip.alias_manager.define_alias(alias, staliases[alias]) 

32 

33 

34def refresh_variables(ip): 

35 db = ip.db 

36 for key in db.keys('autorestore/*'): 

37 # strip autorestore 

38 justkey = os.path.basename(key) 

39 try: 

40 obj = db[key] 

41 except KeyError: 

42 print("Unable to restore variable '%s', ignoring (use %%store -d to forget!)" % justkey) 

43 print("The error was:", sys.exc_info()[0]) 

44 else: 

45 # print("restored",justkey,"=",obj) # dbg 

46 ip.user_ns[justkey] = obj 

47 

48 

49def restore_dhist(ip): 

50 ip.user_ns['_dh'] = ip.db.get('dhist',[]) 

51 

52 

53def restore_data(ip): 

54 refresh_variables(ip) 

55 restore_aliases(ip) 

56 restore_dhist(ip) 

57 

58 

59@magics_class 

60class StoreMagics(Magics): 

61 """Lightweight persistence for python variables. 

62 

63 Provides the %store magic.""" 

64 

65 autorestore = Bool(False, help= 

66 """If True, any %store-d variables will be automatically restored 

67 when IPython starts. 

68 """ 

69 ).tag(config=True) 

70 

71 def __init__(self, shell): 

72 super().__init__(shell=shell) 

73 self.shell.configurables.append(self) 

74 if self.autorestore: 

75 restore_data(self.shell) 

76 

77 @skip_doctest 

78 @line_magic 

79 def store(self, parameter_s=''): 

80 """Lightweight persistence for python variables. 

81 

82 Example:: 

83 

84 In [1]: l = ['hello',10,'world'] 

85 In [2]: %store l 

86 Stored 'l' (list) 

87 In [3]: exit 

88 

89 (IPython session is closed and started again...) 

90 

91 ville@badger:~$ ipython 

92 In [1]: l 

93 NameError: name 'l' is not defined 

94 In [2]: %store -r 

95 In [3]: l 

96 Out[3]: ['hello', 10, 'world'] 

97 

98 Usage: 

99 

100 * ``%store`` - Show list of all variables and their current 

101 values 

102 * ``%store spam bar`` - Store the *current* value of the variables spam 

103 and bar to disk 

104 * ``%store -d spam`` - Remove the variable and its value from storage 

105 * ``%store -z`` - Remove all variables from storage 

106 * ``%store -r`` - Refresh all variables, aliases and directory history 

107 from store (overwrite current vals) 

108 * ``%store -r spam bar`` - Refresh specified variables and aliases from store 

109 (delete current val) 

110 * ``%store foo >a.txt`` - Store value of foo to new file a.txt 

111 * ``%store foo >>a.txt`` - Append value of foo to file a.txt 

112 

113 It should be noted that if you change the value of a variable, you 

114 need to %store it again if you want to persist the new value. 

115 

116 Note also that the variables will need to be pickleable; most basic 

117 python types can be safely %store'd. 

118 

119 Also aliases can be %store'd across sessions. 

120 To remove an alias from the storage, use the %unalias magic. 

121 """ 

122 

123 opts,argsl = self.parse_options(parameter_s,'drz',mode='string') 

124 args = argsl.split() 

125 ip = self.shell 

126 db = ip.db 

127 # delete 

128 if 'd' in opts: 

129 try: 

130 todel = args[0] 

131 except IndexError as e: 

132 raise UsageError('You must provide the variable to forget') from e 

133 else: 

134 try: 

135 del db['autorestore/' + todel] 

136 except BaseException as e: 

137 raise UsageError("Can't delete variable '%s'" % todel) from e 

138 # reset 

139 elif 'z' in opts: 

140 for k in db.keys('autorestore/*'): 

141 del db[k] 

142 

143 elif 'r' in opts: 

144 if args: 

145 for arg in args: 

146 try: 

147 obj = db["autorestore/" + arg] 

148 except KeyError: 

149 try: 

150 restore_aliases(ip, alias=arg) 

151 except KeyError: 

152 print("no stored variable or alias %s" % arg) 

153 else: 

154 ip.user_ns[arg] = obj 

155 else: 

156 restore_data(ip) 

157 

158 # run without arguments -> list variables & values 

159 elif not args: 

160 vars = db.keys('autorestore/*') 

161 vars.sort() 

162 if vars: 

163 size = max(map(len, vars)) 

164 else: 

165 size = 0 

166 

167 print('Stored variables and their in-db values:') 

168 fmt = '%-'+str(size)+'s -> %s' 

169 get = db.get 

170 for var in vars: 

171 justkey = os.path.basename(var) 

172 # print 30 first characters from every var 

173 print(fmt % (justkey, repr(get(var, '<unavailable>'))[:50])) 

174 

175 # default action - store the variable 

176 else: 

177 # %store foo >file.txt or >>file.txt 

178 if len(args) > 1 and args[1].startswith(">"): 

179 fnam = os.path.expanduser(args[1].lstrip(">").lstrip()) 

180 if args[1].startswith(">>"): 

181 fil = open(fnam, "a", encoding="utf-8") 

182 else: 

183 fil = open(fnam, "w", encoding="utf-8") 

184 with fil: 

185 obj = ip.ev(args[0]) 

186 print("Writing '{}' ({}) to file '{}'.".format(args[0], 

187 obj.__class__.__name__, fnam)) 

188 

189 if not isinstance (obj, str): 

190 from pprint import pprint 

191 pprint(obj, fil) 

192 else: 

193 fil.write(obj) 

194 if not obj.endswith('\n'): 

195 fil.write('\n') 

196 

197 return 

198 

199 # %store foo 

200 for arg in args: 

201 try: 

202 obj = ip.user_ns[arg] 

203 except KeyError: 

204 # it might be an alias 

205 name = arg 

206 try: 

207 cmd = ip.alias_manager.retrieve_alias(name) 

208 except ValueError as e: 

209 raise UsageError("Unknown variable '%s'" % name) from e 

210 

211 staliases = db.get('stored_aliases',{}) 

212 staliases[name] = cmd 

213 db['stored_aliases'] = staliases 

214 print("Alias stored: {} ({})".format(name, cmd)) 

215 return 

216 

217 else: 

218 modname = getattr(inspect.getmodule(obj), '__name__', '') 

219 if modname == '__main__': 

220 print(textwrap.dedent("""\ 

221 Warning:{} is {} 

222 Proper storage of interactively declared classes (or instances 

223 of those classes) is not possible! Only instances 

224 of classes in real modules on file system can be %store'd. 

225 """.format(arg, obj) )) 

226 return 

227 #pickled = pickle.dumps(obj) 

228 db[ 'autorestore/' + arg ] = obj 

229 print("Stored '{}' ({})".format(arg, obj.__class__.__name__)) 

230 

231 

232def load_ipython_extension(ip): 

233 """Load the extension in IPython.""" 

234 ip.register_magics(StoreMagics)