From a02911aacbdae6598000fdbd5a08139b24f96b43 2019-10-27 17:47:20 From: Gökcen Eraslan Date: 2019-10-27 17:47:20 Subject: [PATCH] Fix restoring more than 2 variables and add support for storing multiple variables. --- diff --git a/IPython/extensions/storemagic.py b/IPython/extensions/storemagic.py index dae325a..51b79ad 100644 --- a/IPython/extensions/storemagic.py +++ b/IPython/extensions/storemagic.py @@ -20,12 +20,15 @@ from IPython.core.magic import Magics, magics_class, line_magic from traitlets import Bool -def restore_aliases(ip): +def restore_aliases(ip, alias=None): 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) + if alias is None: + for k,v in staliases.items(): + #print "restore alias",k,v # dbg + #self.alias_table[k] = v + ip.alias_manager.define_alias(k,v) + else: + ip.alias_manager.define_alias(alias, staliases[alias]) def refresh_variables(ip): @@ -58,13 +61,13 @@ 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) @@ -94,13 +97,13 @@ class StoreMagics(Magics): * ``%store`` - Show list of all variables and their current values - * ``%store spam`` - Store the *current* value of the variable spam - to disk + * ``%store spam bar`` - Store the *current* value of the variables spam + and bar 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 + * ``%store -r`` - Refresh all variables, aliases and directory history + from store (overwrite current vals) + * ``%store -r spam bar`` - Refresh specified variables and aliases 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 @@ -116,7 +119,7 @@ class StoreMagics(Magics): """ opts,argsl = self.parse_options(parameter_s,'drz',mode='string') - args = argsl.split(None,1) + args = argsl.split() ip = self.shell db = ip.db # delete @@ -141,7 +144,10 @@ class StoreMagics(Magics): try: obj = db['autorestore/' + arg] except KeyError: - print("no stored variable %s" % arg) + try: + restore_aliases(ip, alias=arg) + except KeyError: + print("no stored variable or alias %s" % arg) else: ip.user_ns[arg] = obj else: @@ -189,38 +195,39 @@ class StoreMagics(Magics): return # %store foo - try: - obj = ip.user_ns[args[0]] - except KeyError: - # it might be an alias - name = args[0] + for arg in args: 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) )) + obj = ip.user_ns[arg] + except KeyError: + # it might be an alias + name = arg + 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 - #pickled = pickle.dumps(obj) - db[ 'autorestore/' + args[0] ] = obj - print("Stored '%s' (%s)" % (args[0], obj.__class__.__name__)) + + 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. + """ % (arg, obj) )) + return + #pickled = pickle.dumps(obj) + db[ 'autorestore/' + arg ] = obj + print("Stored '%s' (%s)" % (arg, obj.__class__.__name__)) def load_ipython_extension(ip): """Load the extension in IPython.""" ip.register_magics(StoreMagics) - + diff --git a/IPython/extensions/tests/test_storemagic.py b/IPython/extensions/tests/test_storemagic.py index a7b41ba..6f8371d 100644 --- a/IPython/extensions/tests/test_storemagic.py +++ b/IPython/extensions/tests/test_storemagic.py @@ -10,29 +10,42 @@ def setup_module(): def test_store_restore(): assert 'bar' not in ip.user_ns, "Error: some other test leaked `bar` in user_ns" assert 'foo' not in ip.user_ns, "Error: some other test leaked `foo` in user_ns" + assert 'foobar' not in ip.user_ns, "Error: some other test leaked `foobar` in user_ns" + assert 'foobaz' not in ip.user_ns, "Error: some other test leaked `foobaz` in user_ns" ip.user_ns['foo'] = 78 ip.magic('alias bar echo "hello"') + ip.user_ns['foobar'] = 79 + ip.user_ns['foobaz'] = '80' tmpd = tempfile.mkdtemp() ip.magic('cd ' + tmpd) ip.magic('store foo') ip.magic('store bar') - + ip.magic('store foobar foobaz') + # Check storing nt.assert_equal(ip.db['autorestore/foo'], 78) nt.assert_in('bar', ip.db['stored_aliases']) - + nt.assert_equal(ip.db['autorestore/foobar'], 79) + nt.assert_equal(ip.db['autorestore/foobaz'], '80') + # Remove those items ip.user_ns.pop('foo', None) + ip.user_ns.pop('foobar', None) + ip.user_ns.pop('foobaz', None) ip.alias_manager.undefine_alias('bar') ip.magic('cd -') ip.user_ns['_dh'][:] = [] - + # Check restoring - ip.magic('store -r') + ip.magic('store -r foo bar foobar foobaz') nt.assert_equal(ip.user_ns['foo'], 78) assert ip.alias_manager.is_alias('bar') + nt.assert_equal(ip.user_ns['foobar'], 79) + nt.assert_equal(ip.user_ns['foobaz'], '80') + + ip.magic('store -r') # restores _dh too nt.assert_in(os.path.realpath(tmpd), ip.user_ns['_dh']) - + os.rmdir(tmpd) def test_autorestore():