yaml loader: support env var interpolation in loaded YAML using os.expandvar() for any value ${...} (ukwa/ukwa-pywb#14)

pywb/utils/loaders.py

@@ -11,10 +11,14 @@ import requests
 import yaml
 
 import six
-from six.moves.urllib.parse import unquote_plus, urlsplit
+from six.moves.urllib.parse import unquote_plus, urlsplit, urlencode
 
 import time
 import pkgutil
+import base64
+import yaml
+import cgi
+import re
 
 from io import open, BytesIO
 from warcio.limitreader import LimitReader
@@ -30,7 +34,21 @@ except ImportError:  # pragma: no cover
     s3_avail = False
 
 
-# =================================================================
+# ============================================================================
+def init_yaml_env_vars():
+    env_rx = re.compile(r'\$\{[^}]+\}')
+
+    yaml.add_implicit_resolver('!envvar', env_rx)
+
+    def envvar_constructor(loader, node):
+        value = loader.construct_scalar(node)
+        value = os.path.expandvars(value)
+        return value
+
+    yaml.add_constructor('!envvar', envvar_constructor)
+
+
+# ============================================================================
 def load_py_name(string):
     import importlib
 
@@ -468,3 +486,5 @@ class HMACCookieMaker(object):
 
 # ============================================================================
 BlockLoader.init_default_loaders()
+
+init_yaml_env_vars()

pywb/utils/test/test_loaders.py

@@ -80,6 +80,7 @@ import six
 from six import StringIO
 from io import BytesIO
 import requests
+import yaml
 
 from pywb.utils.loaders import BlockLoader, HMACCookieMaker, to_file_url
 from pywb.utils.loaders import extract_client_cookie
@@ -168,6 +169,28 @@ def test_err_unknown_loader():
 #IOError: No Loader for type: foo
 
 
+
+def test_yaml_resolve_env():
+    os.environ['PYWB_PATH'] = './test'
+    os.environ['PYWB_FOO'] = 'bar'
+
+    config = """\
+collection:
+    coll:
+        index: ${PYWB_PATH}/index
+        archive: ${PYWB_PATH}/archive/${PYWB_FOO}
+        other: ${PYWB_NOT}/archive/${PYWB_FOO}
+"""
+
+    config_data = yaml.load(config)
+
+    assert config_data['collection']['coll']['index'] == './test/index'
+    assert config_data['collection']['coll']['archive'] == './test/archive/bar'
+    assert config_data['collection']['coll']['other'] == '${PYWB_NOT}/archive/bar'
+
+    del os.environ['PYWB_PATH']
+    del os.environ['PYWB_FOO']
+
 def print_str(string):
     return string.decode('utf-8') if six.PY3 else string