storemagic.py
220 lines
| 7.2 KiB
| text/x-python
|
PythonLexer
|
r166 | # -*- coding: utf-8 -*- | ||
""" | ||||
%store magic for lightweight persistence. | ||||
|
r5543 | Stores variables, aliases and macros in IPython's database. | ||
|
r5379 | |||
|
r5543 | To automatically restore stored variables at startup, add this to your | ||
:file:`ipython_config.py` file:: | ||||
|
r5379 | |||
|
r5543 | c.StoreMagic.autorestore = True | ||
|
r166 | """ | ||
|
r7001 | #----------------------------------------------------------------------------- | ||
# Copyright (c) 2012, The IPython Development Team. | ||||
# | ||||
# Distributed under the terms of the Modified BSD License. | ||||
# | ||||
# The full license is in the file COPYING.txt, distributed with this software. | ||||
#----------------------------------------------------------------------------- | ||||
#----------------------------------------------------------------------------- | ||||
# Imports | ||||
#----------------------------------------------------------------------------- | ||||
# Stdlib | ||||
|
r6958 | import inspect, os, sys, textwrap | ||
|
r7001 | # Our own | ||
|
r6958 | from IPython.core.error import UsageError | ||
from IPython.core.fakemodule import FakeModule | ||||
|
r6973 | from IPython.core.magic import Magics, magics_class, line_magic | ||
|
r5544 | from IPython.testing.skipdoctest import skip_doctest | ||
|
r165 | |||
|
r7001 | #----------------------------------------------------------------------------- | ||
# Functions and classes | ||||
#----------------------------------------------------------------------------- | ||||
|
r5541 | |||
|
r5378 | def restore_aliases(ip): | ||
|
r284 | staliases = ip.db.get('stored_aliases', {}) | ||
|
r166 | for k,v in staliases.items(): | ||
#print "restore alias",k,v # dbg | ||||
|
r783 | #self.alias_table[k] = v | ||
|
r5378 | ip.alias_manager.define_alias(k,v) | ||
|
r166 | |||
|
r165 | def refresh_variables(ip): | ||
|
r284 | db = ip.db | ||
|
r165 | 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 | ||||
|
r6935 | print "The error was:", sys.exc_info()[0] | ||
|
r165 | else: | ||
#print "restored",justkey,"=",obj #dbg | ||||
|
r284 | ip.user_ns[justkey] = obj | ||
|
r4872 | |||
|
r165 | |||
|
r713 | def restore_dhist(ip): | ||
|
r5378 | ip.user_ns['_dh'] = ip.db.get('dhist',[]) | ||
|
r4872 | |||
|
r6935 | |||
|
r5378 | def restore_data(ip): | ||
|
r165 | refresh_variables(ip) | ||
|
r5378 | restore_aliases(ip) | ||
restore_dhist(ip) | ||||
|
r165 | |||
|
r6935 | |||
|
r6973 | @magics_class | ||
|
r6935 | class StoreMagics(Magics): | ||
|
r165 | """Lightweight persistence for python variables. | ||
|
r6935 | Provides the %store magic.""" | ||
|
r4872 | |||
|
r6935 | @skip_doctest | ||
@line_magic | ||||
def store(self, parameter_s=''): | ||||
"""Lightweight persistence for python variables. | ||||
|
r4872 | |||
|
r6935 | Example:: | ||
|
r4872 | |||
|
r6935 | In [1]: l = ['hello',10,'world'] | ||
In [2]: %store l | ||||
In [3]: exit | ||||
|
r4872 | |||
|
r6935 | (IPython session is closed and started again...) | ||
|
r4872 | |||
|
r6935 | ville@badger:~$ ipython | ||
In [1]: l | ||||
Out[1]: ['hello', 10, 'world'] | ||||
|
r4872 | |||
|
r6935 | Usage: | ||
|
r4872 | |||
|
r6935 | * ``%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 (delete | ||||
current vals) | ||||
* ``%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 | ||||
|
r4872 | |||
|
r6935 | 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. | ||||
|
r4872 | |||
|
r6935 | Note also that the variables will need to be pickleable; most basic | ||
python types can be safely %store'd. | ||||
|
r165 | |||
|
r6935 | Also aliases can be %store'd across sessions. | ||
""" | ||||
|
r4872 | |||
|
r6935 | opts,argsl = self.parse_options(parameter_s,'drz',mode='string') | ||
args = argsl.split(None,1) | ||||
ip = self.shell | ||||
db = ip.db | ||||
# delete | ||||
|
r7859 | if 'd' in opts: | ||
|
r6935 | try: | ||
todel = args[0] | ||||
except IndexError: | ||||
raise UsageError('You must provide the variable to forget') | ||||
|
r165 | else: | ||
|
r6935 | try: | ||
del db['autorestore/' + todel] | ||||
except: | ||||
raise UsageError("Can't delete variable '%s'" % todel) | ||||
# reset | ||||
|
r7859 | elif 'z' in opts: | ||
|
r6935 | for k in db.keys('autorestore/*'): | ||
del db[k] | ||||
|
r7859 | elif 'r' in opts: | ||
|
r6935 | refresh_variables(ip) | ||
# run without arguments -> list variables & values | ||||
elif not args: | ||||
|
r8002 | vars = db.keys('autorestore/*') | ||
|
r6935 | vars.sort() | ||
if vars: | ||||
|
r6986 | size = max(map(len, vars)) | ||
|
r243 | else: | ||
|
r6935 | 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 | ||||
|
r6986 | print fmt % (justkey, repr(get(var, '<unavailable>'))[:50]) | ||
|
r4872 | |||
|
r6935 | # default action - store the variable | ||
|
r166 | else: | ||
|
r6935 | # %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('>>'): | ||||
|
r6986 | fil = open(fnam, 'a') | ||
|
r6935 | else: | ||
|
r6986 | fil = open(fnam, 'w') | ||
|
r6935 | obj = ip.ev(args[0]) | ||
print "Writing '%s' (%s) to file '%s'." % (args[0], | ||||
obj.__class__.__name__, fnam) | ||||
|
r6986 | if not isinstance (obj, basestring): | ||
|
r6935 | from pprint import pprint | ||
|
r6986 | pprint(obj, fil) | ||
|
r6935 | else: | ||
fil.write(obj) | ||||
if not obj.endswith('\n'): | ||||
fil.write('\n') | ||||
fil.close() | ||||
|
r166 | return | ||
|
r6935 | |||
# %store foo | ||||
try: | ||||
obj = ip.user_ns[args[0]] | ||||
except KeyError: | ||||
# it might be an alias | ||||
# This needs to be refactored to use the new AliasManager stuff. | ||||
|
r8002 | if args[0] in ip.alias_manager: | ||
|
r6935 | name = args[0] | ||
|
r8002 | nargs, cmd = ip.alias_manager.alias_table[ name ] | ||
|
r6935 | staliases = db.get('stored_aliases',{}) | ||
staliases[ name ] = cmd | ||||
db['stored_aliases'] = staliases | ||||
print "Alias stored: %s (%s)" % (name, cmd) | ||||
return | ||||
else: | ||||
raise UsageError("Unknown variable '%s'" % args[0]) | ||||
else: | ||||
if isinstance(inspect.getmodule(obj), FakeModule): | ||||
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) | ||||
|
r8002 | db[ 'autorestore/' + args[0] ] = obj | ||
|
r6935 | print "Stored '%s' (%s)" % (args[0], obj.__class__.__name__) | ||
|
r165 | |||
|
r5541 | |||
_loaded = False | ||||
|
r8197 | |||
|
r5378 | def load_ipython_extension(ip): | ||
|
r5541 | """Load the extension in IPython.""" | ||
global _loaded | ||||
if not _loaded: | ||||
|
r8197 | ip.register_magics(StoreMagics) | ||
|
r5541 | _loaded = True | ||