From 2de04994df5967d294cb4263aa467cc76e528bed Mon Sep 17 00:00:00 2001
From: Steven Loria <sloria1@gmail.com>
Date: Sat, 12 May 2018 12:30:00 -0400
Subject: [PATCH] Support setting envvars with 'export' commands

closes #32
---
 CHANGELOG.rst                 |  2 ++
 doitlive/cli.py               | 12 +++++++++++-
 tests/sessions/export.session |  2 ++
 tests/test_doitlive.py        |  9 +++++++++
 4 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 tests/sessions/export.session

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 7dcf789..5585357 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -8,6 +8,8 @@ Features:
 
 * Add shell completion for bash, zsh, and fish (:issue:`3`).
 * Add "Did you mean" suggestions.
+* Support setting environment variables with ``export`` commands
+  (:issue:`32`). Thanks :user:`asmacdo` for the suggestion.
 
 Bug fixes:
 
diff --git a/doitlive/cli.py b/doitlive/cli.py
index 9299772..ade9574 100755
--- a/doitlive/cli.py
+++ b/doitlive/cli.py
@@ -2,6 +2,7 @@
 import functools
 import os
 import re
+import shlex
 import sys
 import textwrap
 from codecs import open
@@ -11,6 +12,7 @@
 from click import secho, style
 from click_didyoumean import DYMGroup
 
+from doitlive.__version__ import __version__
 from doitlive.compat import ensure_utf8
 from doitlive.exceptions import SessionError
 from doitlive.keyboard import (RETURNS, magicrun, magictype, run_command,
@@ -18,7 +20,6 @@
 from doitlive.python_consoles import PythonRecorderConsole, start_python_player
 from doitlive.styling import THEMES, echo, echo_prompt, format_prompt
 from doitlive.termutils import get_default_shell
-from doitlive.__version__ import __version__
 
 env = os.environ
 click_completion.init()
@@ -120,6 +121,7 @@ def run(commands, shell=None, prompt_template='default', speed=1,
     i = 0
     while i < len(commands):
         command = commands[i].strip()
+        command_as_list = shlex.split(command)
         i += 1
         if not command:
             continue
@@ -135,6 +137,14 @@ def run(commands, shell=None, prompt_template='default', speed=1,
                 comment = command.lstrip("#")
                 secho(comment, fg='yellow', bold=True)
             continue
+        # Handle 'export' commands by adding envvars to SessionState
+        elif command_as_list and command_as_list[0] == 'export':
+            magictype(command,
+                      prompt_template=state['prompt_template'],
+                      speed=state['speed'])
+            for envvar in command_as_list[1:]:
+                state.add_envvar(envvar)
+        # Handle ```python and ```ipython by running "player" shells
         elif shell_match:
             shell_name = shell_match.groups()[0].strip()
             py_commands = []
diff --git a/tests/sessions/export.session b/tests/sessions/export.session
new file mode 100644
index 0000000..918189e
--- /dev/null
+++ b/tests/sessions/export.session
@@ -0,0 +1,2 @@
+export NAME=Steve
+echo 'Hello' $NAME
diff --git a/tests/test_doitlive.py b/tests/test_doitlive.py
index 5362859..e518422 100644
--- a/tests/test_doitlive.py
+++ b/tests/test_doitlive.py
@@ -155,6 +155,15 @@ def test_unset_envvar(self, runner):
         result = run_session(runner, 'unset.session', user_input)
         assert 'fortytwo' not in result.output
 
+    def test_export(self, runner):
+        user_input = ''.join([
+            random_string(len('export NAME=Steve')),
+            '\n',
+            random_string(len("echo 'Hello' $NAME"))
+        ])
+        result = run_session(runner, 'export.session', user_input)
+        assert 'Hello Steve' in result.output
+
 
 def test_themes_list(runner):
     result1 = runner.invoke(cli, ['themes'])