aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/python/ipython/py2/IPython/extensions/storemagic.py
blob: 34371a30f231d7388d2d370a6a23b91bcc991279 (plain) (blame)
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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
# -*- coding: utf-8 -*- 
""" 
%store magic for lightweight persistence. 
 
Stores variables, aliases and macros in IPython's database. 
 
To automatically restore stored variables at startup, add this to your 
:file:`ipython_config.py` file:: 
 
  c.StoreMagics.autorestore = True 
""" 
from __future__ import print_function 
 
# Copyright (c) IPython Development Team.
# Distributed under the terms of the Modified BSD License.
 
import inspect, os, sys, textwrap 
 
from IPython.core.error import UsageError 
from IPython.core.magic import Magics, magics_class, line_magic 
from traitlets import Bool 
from IPython.utils.py3compat import string_types 
 
 
def restore_aliases(ip): 
    staliases = ip.db.get('stored_aliases', {}) 
    for k,v in staliases.items(): 
        #print "restore alias",k,v # dbg 
        #self.alias_table[k] = v 
        ip.alias_manager.define_alias(k,v) 
 
 
def refresh_variables(ip): 
    db = ip.db 
    for key in db.keys('autorestore/*'): 
        # strip autorestore 
        justkey = os.path.basename(key) 
        try: 
            obj = db[key] 
        except KeyError: 
            print("Unable to restore variable '%s', ignoring (use %%store -d to forget!)" % justkey) 
            print("The error was:", sys.exc_info()[0]) 
        else: 
            #print "restored",justkey,"=",obj #dbg 
            ip.user_ns[justkey] = obj 
 
 
def restore_dhist(ip): 
    ip.user_ns['_dh'] = ip.db.get('dhist',[]) 
 
 
def restore_data(ip): 
    refresh_variables(ip) 
    restore_aliases(ip) 
    restore_dhist(ip) 
 
 
@magics_class 
class StoreMagics(Magics): 
    """Lightweight persistence for python variables. 
 
    Provides the %store magic.""" 
     
    autorestore = Bool(False, help=
        """If True, any %store-d variables will be automatically restored 
        when IPython starts. 
        """ 
    ).tag(config=True)
     
    def __init__(self, shell): 
        super(StoreMagics, self).__init__(shell=shell) 
        self.shell.configurables.append(self) 
        if self.autorestore: 
            restore_data(self.shell) 
 
    @line_magic 
    def store(self, parameter_s=''): 
        """Lightweight persistence for python variables. 
 
        Example:: 
 
          In [1]: l = ['hello',10,'world'] 
          In [2]: %store l 
          In [3]: exit 
 
          (IPython session is closed and started again...) 
 
          ville@badger:~$ ipython 
          In [1]: l 
          NameError: name 'l' is not defined 
          In [2]: %store -r 
          In [3]: l 
          Out[3]: ['hello', 10, 'world'] 
 
        Usage: 
 
        * ``%store``          - Show list of all variables and their current 
                                values 
        * ``%store spam``     - Store the *current* value of the variable spam 
                                to disk 
        * ``%store -d spam``  - Remove the variable and its value from storage 
        * ``%store -z``       - Remove all variables from storage 
        * ``%store -r``       - Refresh all variables from store (overwrite 
                                current vals) 
        * ``%store -r spam bar`` - Refresh specified variables from store 
                                   (delete current val) 
        * ``%store foo >a.txt``  - Store value of foo to new file a.txt 
        * ``%store foo >>a.txt`` - Append value of foo to file a.txt 
 
        It should be noted that if you change the value of a variable, you 
        need to %store it again if you want to persist the new value. 
 
        Note also that the variables will need to be pickleable; most basic 
        python types can be safely %store'd. 
 
        Also aliases can be %store'd across sessions. 
        """ 
 
        opts,argsl = self.parse_options(parameter_s,'drz',mode='string') 
        args = argsl.split(None,1) 
        ip = self.shell 
        db = ip.db 
        # delete 
        if 'd' in opts: 
            try: 
                todel = args[0] 
            except IndexError: 
                raise UsageError('You must provide the variable to forget') 
            else: 
                try: 
                    del db['autorestore/' + todel] 
                except: 
                    raise UsageError("Can't delete variable '%s'" % todel) 
        # reset 
        elif 'z' in opts: 
            for k in db.keys('autorestore/*'): 
                del db[k] 
 
        elif 'r' in opts: 
            if args: 
                for arg in args: 
                    try: 
                        obj = db['autorestore/' + arg] 
                    except KeyError: 
                        print("no stored variable %s" % arg) 
                    else: 
                        ip.user_ns[arg] = obj 
            else: 
                restore_data(ip) 
 
        # run without arguments -> list variables & values 
        elif not args: 
            vars = db.keys('autorestore/*') 
            vars.sort() 
            if vars: 
                size = max(map(len, vars)) 
            else: 
                size = 0 
 
            print('Stored variables and their in-db values:') 
            fmt = '%-'+str(size)+'s -> %s' 
            get = db.get 
            for var in vars: 
                justkey = os.path.basename(var) 
                # print 30 first characters from every var 
                print(fmt % (justkey, repr(get(var, '<unavailable>'))[:50])) 
 
        # default action - store the variable 
        else: 
            # %store foo >file.txt or >>file.txt 
            if len(args) > 1 and args[1].startswith('>'): 
                fnam = os.path.expanduser(args[1].lstrip('>').lstrip()) 
                if args[1].startswith('>>'): 
                    fil = open(fnam, 'a') 
                else: 
                    fil = open(fnam, 'w') 
                obj = ip.ev(args[0]) 
                print("Writing '%s' (%s) to file '%s'." % (args[0], 
                  obj.__class__.__name__, fnam)) 
 
 
                if not isinstance (obj, string_types): 
                    from pprint import pprint 
                    pprint(obj, fil) 
                else: 
                    fil.write(obj) 
                    if not obj.endswith('\n'): 
                        fil.write('\n') 
 
                fil.close() 
                return 
 
            # %store foo 
            try: 
                obj = ip.user_ns[args[0]] 
            except KeyError: 
                # it might be an alias 
                name = args[0] 
                try: 
                    cmd = ip.alias_manager.retrieve_alias(name) 
                except ValueError: 
                    raise UsageError("Unknown variable '%s'" % name) 
                 
                staliases = db.get('stored_aliases',{}) 
                staliases[name] = cmd 
                db['stored_aliases'] = staliases 
                print("Alias stored: %s (%s)" % (name, cmd)) 
                return 
 
            else: 
                modname = getattr(inspect.getmodule(obj), '__name__', '') 
                if modname == '__main__': 
                    print(textwrap.dedent("""\ 
                    Warning:%s is %s 
                    Proper storage of interactively declared classes (or instances 
                    of those classes) is not possible! Only instances 
                    of classes in real modules on file system can be %%store'd. 
                    """ % (args[0], obj) )) 
                    return 
                #pickled = pickle.dumps(obj) 
                db[ 'autorestore/' + args[0] ] = obj 
                print("Stored '%s' (%s)" % (args[0], obj.__class__.__name__)) 
 
 
def load_ipython_extension(ip): 
    """Load the extension in IPython.""" 
    ip.register_magics(StoreMagics)