From 8b9dac50eca8d2737e377c225fdd5e4c4dc948c8 Mon Sep 17 00:00:00 2001
From: cwandrews <cwandrews.oh@gmail.com>
Date: Mon, 18 Sep 2017 15:24:27 -0400
Subject: [PATCH] Added optional passing of one or multiple environment
 variables via --env flag to the bonobo cli.

---
 bonobo/commands/run.py       | 12 +++++++++++-
 tests/test_commands.py       | 27 +++++++++++++++++++++++++++
 tests/util/get_passed_env.py | 22 ++++++++++++++++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 tests/util/get_passed_env.py

diff --git a/bonobo/commands/run.py b/bonobo/commands/run.py
index fb93e77..c9bb39e 100644
--- a/bonobo/commands/run.py
+++ b/bonobo/commands/run.py
@@ -40,7 +40,10 @@ def _install_requirements(requirements):
     importlib.reload(site)
 
 
-def execute(filename, module, install=False, quiet=False, verbose=False):
+def execute(filename, module, install=False, quiet=False, verbose=False,
+            env=None):
+    import re
+
     import runpy
     from bonobo import Graph, run, settings
 
@@ -50,6 +53,12 @@ def execute(filename, module, install=False, quiet=False, verbose=False):
     if verbose:
         settings.DEBUG.set(True)
 
+    if env:
+        quote_killer = re.compile('["\']')
+        for e in env:
+            var_name, var_value = e.split('=')
+            os.environ[var_name] = quote_killer.sub('', var_value)
+
     if filename:
         if os.path.isdir(filename):
             if install:
@@ -106,4 +115,5 @@ def register(parser):
     verbosity_group.add_argument('--quiet', '-q', action='store_true')
     verbosity_group.add_argument('--verbose', '-v', action='store_true')
     parser.add_argument('--install', '-I', action='store_true')
+    parser.add_argument('--env', '-e', action='append')
     return execute
diff --git a/tests/test_commands.py b/tests/test_commands.py
index daf245f..cff9e38 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -3,6 +3,7 @@ import runpy
 import sys
 from unittest.mock import patch
 
+import pathlib
 import pkg_resources
 import pytest
 
@@ -96,3 +97,29 @@ def test_version(runner, capsys):
     out = out.strip()
     assert out.startswith('bonobo ')
     assert __version__ in out
+
+
+@all_runners
+def test_run_with_env(runner, capsys):
+    runner('run', '--quiet',
+           str(pathlib.Path(os.path.dirname(__file__),
+                            'util', 'get_passed_env.py')),
+           '--env', 'ENV_TEST_NUMBER=123', '--env', 'ENV_TEST_USER=cwandrews',
+           '--env', "ENV_TEST_STRING='my_test_string'")
+    out, err = capsys.readouterr()
+    out = out.split('\n')
+    assert out[0] == 'cwandrews'
+    assert out[1] == '123'
+    assert out[2] == 'my_test_string'
+
+
+@all_runners
+def test_run_module_with_env(runner, capsys):
+    runner('run', '--quiet', '-m', 'tests.util.get_passed_env',
+           '--env', 'ENV_TEST_NUMBER=123', '--env', 'ENV_TEST_USER=cwandrews',
+           '--env', "ENV_TEST_STRING='my_test_string'")
+    out, err = capsys.readouterr()
+    out = out.split('\n')
+    assert out[0] == 'cwandrews'
+    assert out[1] == '123'
+    assert out[2] == 'my_test_string'
diff --git a/tests/util/get_passed_env.py b/tests/util/get_passed_env.py
new file mode 100644
index 0000000..d9c4ba6
--- /dev/null
+++ b/tests/util/get_passed_env.py
@@ -0,0 +1,22 @@
+import os
+
+from bonobo import Graph
+
+
+def extract():
+    env_test_user = os.getenv('ENV_TEST_USER')
+    env_test_number = os.getenv('ENV_TEST_NUMBER')
+    env_test_string = os.getenv('ENV_TEST_STRING')
+    return env_test_user, env_test_number, env_test_string
+
+
+def load(s: str):
+    print(s)
+
+
+graph = Graph(extract, load)
+
+if __name__ == '__main__':
+    from bonobo import run
+
+    run(graph)