Show More
@@ -371,14 +371,52 b' class OSMagics(Magics):' | |||
|
371 | 371 | if not 'q' in opts and self.shell.user_ns['_dh']: |
|
372 | 372 | print(self.shell.user_ns['_dh'][-1]) |
|
373 | 373 | |
|
374 | ||
|
375 | 374 | @line_magic |
|
376 | 375 | def env(self, parameter_s=''): |
|
377 | 376 | """List environment variables.""" |
|
378 | ||
|
377 | if parameter_s.strip(): | |
|
378 | split = '=' if '=' in parameter_s else ' ' | |
|
379 | bits = parameter_s.split(split) | |
|
380 | if len(bits) == 1: | |
|
381 | key = parameter_s.strip() | |
|
382 | if key in os.environ: | |
|
383 | return os.environ[key] | |
|
384 | else: | |
|
385 | err = "Environment does not have key: {0}".format(key) | |
|
386 | raise UsageError(err) | |
|
387 | if len(bits) > 1: | |
|
388 | return self.set_env(parameter_s) | |
|
379 | 389 | return dict(os.environ) |
|
380 | 390 | |
|
381 | 391 | @line_magic |
|
392 | def set_env(self, parameter_s): | |
|
393 | """Set environment variables. Assumptions are that either "val" is a | |
|
394 | name in the user namespace, or val is something that evaluates to a | |
|
395 | string. | |
|
396 | ||
|
397 | Usage:\\ | |
|
398 | %set_env var val | |
|
399 | """ | |
|
400 | split = '=' if '=' in parameter_s else ' ' | |
|
401 | bits = parameter_s.split(split, 1) | |
|
402 | if not parameter_s.strip() or len(bits)<2: | |
|
403 | raise UsageError("usage is 'set_env var=val'") | |
|
404 | var = bits[0].strip() | |
|
405 | val = bits[1].strip() | |
|
406 | if re.match(r'.*\s.*', var): | |
|
407 | # an environment variable with whitespace is almost certainly | |
|
408 | # not what the user intended. what's more likely is the wrong | |
|
409 | # split was chosen, ie for "set_env cmd_args A=B", we chose | |
|
410 | # '=' for the split and should have chosen ' '. to get around | |
|
411 | # this, users should just assign directly to os.environ or use | |
|
412 | # standard magic {var} expansion. | |
|
413 | err = "refusing to set env var with whitespace: '{0}'" | |
|
414 | err = err.format(val) | |
|
415 | raise UsageError(err) | |
|
416 | os.environ[var] = val | |
|
417 | print('env: {0}={1}'.format(var,val)) | |
|
418 | ||
|
419 | @line_magic | |
|
382 | 420 | def pushd(self, parameter_s=''): |
|
383 | 421 | """Place the current dir on stack and change directory. |
|
384 | 422 |
@@ -19,6 +19,7 b' except ImportError:' | |||
|
19 | 19 | import nose.tools as nt |
|
20 | 20 | |
|
21 | 21 | from IPython.core import magic |
|
22 | from IPython.core.error import UsageError | |
|
22 | 23 | from IPython.core.magic import (Magics, magics_class, line_magic, |
|
23 | 24 | cell_magic, line_cell_magic, |
|
24 | 25 | register_line_magic, register_cell_magic, |
@@ -626,9 +627,35 b' class NotebookExportMagicTests(TestCase):' | |||
|
626 | 627 | _ip.magic("notebook -e %s" % outfile) |
|
627 | 628 | |
|
628 | 629 | |
|
629 | def test_env(): | |
|
630 | env = _ip.magic("env") | |
|
631 | assert isinstance(env, dict), type(env) | |
|
630 | class TestEnv(TestCase): | |
|
631 | ||
|
632 | def test_env(self): | |
|
633 | env = _ip.magic("env") | |
|
634 | self.assertTrue(isinstance(env, dict)) | |
|
635 | ||
|
636 | def test_env_get_set_simple(self): | |
|
637 | env = _ip.magic("env var val1") | |
|
638 | self.assertEqual(env, None) | |
|
639 | self.assertEqual(os.environ['var'], 'val1') | |
|
640 | self.assertEqual(_ip.magic("env var"), 'val1') | |
|
641 | env = _ip.magic("env var=val2") | |
|
642 | self.assertEqual(env, None) | |
|
643 | self.assertEqual(os.environ['var'], 'val2') | |
|
644 | ||
|
645 | def test_env_get_set_complex(self): | |
|
646 | env = _ip.magic("env var 'val1 '' 'val2") | |
|
647 | self.assertEqual(env, None) | |
|
648 | self.assertEqual(os.environ['var'], "'val1 '' 'val2") | |
|
649 | self.assertEqual(_ip.magic("env var"), "'val1 '' 'val2") | |
|
650 | env = _ip.magic('env var=val2 val3="val4') | |
|
651 | self.assertEqual(env, None) | |
|
652 | self.assertEqual(os.environ['var'], 'val2 val3="val4') | |
|
653 | ||
|
654 | def test_env_set_bad_input(self): | |
|
655 | self.assertRaises(UsageError, lambda: _ip.magic("set_env var")) | |
|
656 | ||
|
657 | def test_env_set_whitespace(self): | |
|
658 | self.assertRaises(UsageError, lambda: _ip.magic("env var A=B")) | |
|
632 | 659 | |
|
633 | 660 | |
|
634 | 661 | class CellMagicTestCase(TestCase): |
General Comments 0
You need to be logged in to leave comments.
Login now