~ellisonbg/ipython/bugfixes0411409

« back to all changes in this revision

Viewing changes to IPython/Extensions/pspersistence.py

  • Committer: ville
  • Date: 2008-02-16 09:50:47 UTC
  • mto: (0.12.1 ipython_main)
  • mto: This revision was merged to the branch mainline in revision 990.
  • Revision ID: ville@ville-pc-20080216095047-500x6dluki1iz40o
initialization (no svn history)

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# -*- coding: utf-8 -*-
 
2
"""
 
3
%store magic for lightweight persistence.
 
4
 
 
5
Stores variables, aliases etc. in PickleShare database.
 
6
 
 
7
$Id: iplib.py 1107 2006-01-30 19:02:20Z vivainio $
 
8
"""
 
9
 
 
10
import IPython.ipapi
 
11
ip = IPython.ipapi.get()
 
12
 
 
13
import pickleshare
 
14
 
 
15
import inspect,pickle,os,sys,textwrap
 
16
from IPython.FakeModule import FakeModule
 
17
 
 
18
def restore_aliases(self):
 
19
    ip = self.getapi()
 
20
    staliases = ip.db.get('stored_aliases', {})
 
21
    for k,v in staliases.items():
 
22
        #print "restore alias",k,v # dbg
 
23
        #self.alias_table[k] = v
 
24
        ip.defalias(k,v)
 
25
 
 
26
 
 
27
def refresh_variables(ip):
 
28
    db = ip.db
 
29
    for key in db.keys('autorestore/*'):
 
30
        # strip autorestore
 
31
        justkey = os.path.basename(key)
 
32
        try:
 
33
            obj = db[key]
 
34
        except KeyError:
 
35
            print "Unable to restore variable '%s', ignoring (use %%store -d to forget!)" % justkey
 
36
            print "The error was:",sys.exc_info()[0]
 
37
        else:
 
38
            #print "restored",justkey,"=",obj #dbg
 
39
            ip.user_ns[justkey] = obj
 
40
    
 
41
 
 
42
def restore_dhist(ip):
 
43
    db = ip.db
 
44
    ip.user_ns['_dh'] = db.get('dhist',[])
 
45
    
 
46
def restore_data(self):
 
47
    ip = self.getapi()
 
48
    refresh_variables(ip)
 
49
    restore_aliases(self)
 
50
    restore_dhist(self)
 
51
    raise IPython.ipapi.TryNext
 
52
    
 
53
ip.set_hook('late_startup_hook', restore_data)
 
54
 
 
55
def magic_store(self, parameter_s=''):
 
56
    """Lightweight persistence for python variables.
 
57
 
 
58
    Example:
 
59
    
 
60
    ville@badger[~]|1> A = ['hello',10,'world']\\
 
61
    ville@badger[~]|2> %store A\\
 
62
    ville@badger[~]|3> Exit
 
63
    
 
64
    (IPython session is closed and started again...)
 
65
    
 
66
    ville@badger:~$ ipython -p pysh\\
 
67
    ville@badger[~]|1> print A
 
68
    
 
69
    ['hello', 10, 'world']
 
70
    
 
71
    Usage:
 
72
    
 
73
    %store          - Show list of all variables and their current values\\
 
74
    %store <var>    - Store the *current* value of the variable to disk\\
 
75
    %store -d <var> - Remove the variable and its value from storage\\
 
76
    %store -z       - Remove all variables from storage\\
 
77
    %store -r       - Refresh all variables from store (delete current vals)\\
 
78
    %store foo >a.txt  - Store value of foo to new file a.txt\\
 
79
    %store foo >>a.txt - Append value of foo to file a.txt\\   
 
80
    
 
81
    It should be noted that if you change the value of a variable, you
 
82
    need to %store it again if you want to persist the new value.
 
83
    
 
84
    Note also that the variables will need to be pickleable; most basic
 
85
    python types can be safely %stored.
 
86
    
 
87
    Also aliases can be %store'd across sessions.
 
88
    """
 
89
    
 
90
    opts,argsl = self.parse_options(parameter_s,'drz',mode='string')
 
91
    args = argsl.split(None,1)
 
92
    ip = self.getapi()
 
93
    db = ip.db
 
94
    # delete
 
95
    if opts.has_key('d'):
 
96
        try:
 
97
            todel = args[0]
 
98
        except IndexError:
 
99
            error('You must provide the variable to forget')
 
100
        else:
 
101
            try:
 
102
                del db['autorestore/' + todel]
 
103
            except:
 
104
                error("Can't delete variable '%s'" % todel)
 
105
    # reset
 
106
    elif opts.has_key('z'):
 
107
        for k in db.keys('autorestore/*'):
 
108
            del db[k]
 
109
 
 
110
    elif opts.has_key('r'):
 
111
        refresh_variables(ip)
 
112
 
 
113
    
 
114
    # run without arguments -> list variables & values
 
115
    elif not args:
 
116
        vars = self.db.keys('autorestore/*')
 
117
        vars.sort()            
 
118
        if vars:
 
119
            size = max(map(len,vars))
 
120
        else:
 
121
            size = 0
 
122
            
 
123
        print 'Stored variables and their in-db values:'
 
124
        fmt = '%-'+str(size)+'s -> %s'
 
125
        get = db.get
 
126
        for var in vars:
 
127
            justkey = os.path.basename(var)
 
128
            # print 30 first characters from every var
 
129
            print fmt % (justkey,repr(get(var,'<unavailable>'))[:50])
 
130
    
 
131
    # default action - store the variable
 
132
    else:
 
133
        # %store foo >file.txt or >>file.txt
 
134
        if len(args) > 1 and args[1].startswith('>'):
 
135
            fnam = os.path.expanduser(args[1].lstrip('>').lstrip())
 
136
            if args[1].startswith('>>'):
 
137
                fil = open(fnam,'a')
 
138
            else:
 
139
                fil = open(fnam,'w')
 
140
            obj = ip.ev(args[0])
 
141
            print "Writing '%s' (%s) to file '%s'." % (args[0],
 
142
              obj.__class__.__name__, fnam)
 
143
 
 
144
            
 
145
            if not isinstance (obj,basestring):
 
146
                from pprint import pprint
 
147
                pprint(obj,fil)
 
148
            else:
 
149
                fil.write(obj)
 
150
                if not obj.endswith('\n'):
 
151
                    fil.write('\n')
 
152
            
 
153
            fil.close()
 
154
            return
 
155
        
 
156
        # %store foo
 
157
        try:
 
158
            obj = ip.user_ns[args[0]]
 
159
        except KeyError:
 
160
            # it might be an alias
 
161
            if args[0] in self.alias_table:
 
162
                staliases = db.get('stored_aliases',{})
 
163
                staliases[ args[0] ] = self.alias_table[ args[0] ]
 
164
                db['stored_aliases'] = staliases                
 
165
                print "Alias stored:", args[0], self.alias_table[ args[0] ]
 
166
                return
 
167
            else:
 
168
                print "Error: unknown variable '%s'" % args[0]
 
169
            
 
170
        else:
 
171
            if isinstance(inspect.getmodule(obj), FakeModule):
 
172
                print textwrap.dedent("""\
 
173
                Warning:%s is %s 
 
174
                Proper storage of interactively declared classes (or instances
 
175
                of those classes) is not possible! Only instances
 
176
                of classes in real modules on file system can be %%store'd.
 
177
                """ % (args[0], obj) ) 
 
178
                return
 
179
            #pickled = pickle.dumps(obj)
 
180
            self.db[ 'autorestore/' + args[0] ] = obj
 
181
            print "Stored '%s' (%s)" % (args[0], obj.__class__.__name__)
 
182
 
 
183
ip.expose_magic('store',magic_store)