From f5839c2fc491f2e431a51273f020f8a8001e0e32 2014-11-18 13:41:15
From: Matthias Bussonnier <bussonniermatthias@gmail.com>
Date: 2014-11-18 13:41:15
Subject: [PATCH] Merge pull request #6953 from mattvonrocketstein/issue3992

[issue3992] set/get use cases for 'env' and new 'set_env' magic
---

diff --git a/IPython/core/magics/osm.py b/IPython/core/magics/osm.py
index ccb2105..729fcb7 100644
--- a/IPython/core/magics/osm.py
+++ b/IPython/core/magics/osm.py
@@ -371,14 +371,52 @@ class OSMagics(Magics):
         if not 'q' in opts and self.shell.user_ns['_dh']:
             print(self.shell.user_ns['_dh'][-1])
 
-
     @line_magic
     def env(self, parameter_s=''):
         """List environment variables."""
-
+        if parameter_s.strip():
+            split = '=' if '=' in parameter_s else ' '
+            bits = parameter_s.split(split)
+            if len(bits) == 1:
+                key = parameter_s.strip()
+                if key in os.environ:
+                    return os.environ[key]
+                else:
+                    err = "Environment does not have key: {0}".format(key)
+                    raise UsageError(err)
+            if len(bits) > 1:
+                return self.set_env(parameter_s)
         return dict(os.environ)
 
     @line_magic
+    def set_env(self, parameter_s):
+        """Set environment variables.  Assumptions are that either "val" is a
+        name in the user namespace, or val is something that evaluates to a
+        string.
+
+        Usage:\\
+        %set_env var val
+        """
+        split = '=' if '=' in parameter_s else ' '
+        bits = parameter_s.split(split, 1)
+        if not parameter_s.strip() or len(bits)<2:
+            raise UsageError("usage is 'set_env var=val'")
+        var = bits[0].strip()
+        val = bits[1].strip()
+        if re.match(r'.*\s.*', var):
+            # an environment variable with whitespace is almost certainly
+            # not what the user intended.  what's more likely is the wrong
+            # split was chosen, ie for "set_env cmd_args A=B", we chose
+            # '=' for the split and should have chosen ' '.  to get around
+            # this, users should just assign directly to os.environ or use
+            # standard magic {var} expansion.
+            err = "refusing to set env var with whitespace: '{0}'"
+            err = err.format(val)
+            raise UsageError(err)
+        os.environ[var] = val
+        print('env: {0}={1}'.format(var,val))
+
+    @line_magic
     def pushd(self, parameter_s=''):
         """Place the current dir on stack and change directory.
 
diff --git a/IPython/core/tests/test_magic.py b/IPython/core/tests/test_magic.py
index acd0d3d..a8b9cc6 100644
--- a/IPython/core/tests/test_magic.py
+++ b/IPython/core/tests/test_magic.py
@@ -19,6 +19,7 @@ except ImportError:
 import nose.tools as nt
 
 from IPython.core import magic
+from IPython.core.error import UsageError
 from IPython.core.magic import (Magics, magics_class, line_magic,
                                 cell_magic, line_cell_magic,
                                 register_line_magic, register_cell_magic,
@@ -626,9 +627,35 @@ class NotebookExportMagicTests(TestCase):
             _ip.magic("notebook -e %s" % outfile)
 
 
-def test_env():
-    env = _ip.magic("env")
-    assert isinstance(env, dict), type(env)
+class TestEnv(TestCase):
+
+    def test_env(self):
+        env = _ip.magic("env")
+        self.assertTrue(isinstance(env, dict))
+
+    def test_env_get_set_simple(self):
+        env = _ip.magic("env var val1")
+        self.assertEqual(env, None)
+        self.assertEqual(os.environ['var'], 'val1')
+        self.assertEqual(_ip.magic("env var"), 'val1')
+        env = _ip.magic("env var=val2")
+        self.assertEqual(env, None)
+        self.assertEqual(os.environ['var'], 'val2')
+
+    def test_env_get_set_complex(self):
+        env = _ip.magic("env var 'val1 '' 'val2")
+        self.assertEqual(env, None)
+        self.assertEqual(os.environ['var'], "'val1 '' 'val2")
+        self.assertEqual(_ip.magic("env var"), "'val1 '' 'val2")
+        env = _ip.magic('env var=val2 val3="val4')
+        self.assertEqual(env, None)
+        self.assertEqual(os.environ['var'], 'val2 val3="val4')
+
+    def test_env_set_bad_input(self):
+        self.assertRaises(UsageError, lambda: _ip.magic("set_env var"))
+
+    def test_env_set_whitespace(self):
+        self.assertRaises(UsageError, lambda: _ip.magic("env var A=B"))
 
 
 class CellMagicTestCase(TestCase):