Repository: dask/dask
Pull request: add expand_environment_variables to dask.config
closes #3879

- [x] Tests added / passed
- [x] Passes `flake8 dask`

cc @jacobtomlinson and @mrocklin 

Problem statement:
utility to expand environment variables in dask config
It would be useful if dask could (optionally) expand environment variables in its config module. This may not always be useful but @mrocklin and I ran into a case today where it is the logical solution. Consider the following dask configuration file.

```yaml
# dask config file...
kubernetes:
  worker-template:
    metadata:
    spec:
      containers:
      - args:
          - dask-worker
          - --nthreads
          - '2'
        image: ${JUPYTER_IMAGE_SPEC}
        name: dask-worker
```

Origianally, I had intended to set the image variable by setting  `${DASK_KUBERNETES__WORKER_TEMPLATE__SPEC__IMAGE}`. However, in this case, `containers` is a list and the current environment variable approach is not able to handle indexing in the 0th position of the list. A simple python snippit like this does the required indexing but is clearly not general enough to be useful beyond my immediate application:

```Python
dask.config.config['kubernetes']['worker-template']['spec']['containers'][0]['image'] = os.environ['JUPYTER_IMAGE_SPEC']
```

------

At the moment, this feature seems most useful to dask-kubernetes. I know we don't want to expand environment variables in dask-jobqueue. So this seems like something that we would want to users/downstream packages to be able to opt into. I'm thinking a utility in `dask.config` to optionally expand environment variables would useful.

Hints from issue discussion:
Right, so this would be a utility function in `dask/config.py` and then a use of that utility function in dask-kubernetes?
Exactly. We should ping @jacobtomlinson for comment as well.
We do currently use this in [a few places](https://github.com/dask/dask-kubernetes/blob/master/dask_kubernetes/core.py#L190) in dask-kubernetes to expand some config items.

Are you suggesting that we override `dask.config.get` to take a boolean kwarg to enable this instead of formatting manually?
I think that, to start, it would just be a separate function.  Perhaps something like the following:

```python
spec = dask.config.get('kubernetes.worker-template')
spec = expand_environment_variables(spec)
```

This is valuable in this case because the environment variable is buried somewhat deeply.   
Do you feel this adds much value over

```
spec = dask.config.get('kubernetes.worker-template')
spec = spec.format(**os.environ)
```
Recall that spec in this case may be a nested dictionary including the full
worker template kubernetes specification, not just a string pointing to
another location on disk.

In @jhamman's case he wants to include an environment variable nested well
within that dictionary for the container's image and wants that environment
variable to be expanded

On Mon, Aug 20, 2018 at 8:28 AM, Jacob Tomlinson <notifications@github.com>
wrote:

> Do you feel this adds much value over
>
> spec = dask.config.get('kubernetes.worker-template')
> spec = spec.format(**os.environ)
>
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/3879#issuecomment-414299603>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszNWV3ETEdav0xFV18YBp_yyF7MAjks5uSqtbgaJpZM4V_K5v>
> .
>

Ah that makes more sense. So `expand_environment_variables` would iterate over the dictionary expanding all the variables.

Sounds good.

Commit message: add expand_environment_variables to dask.config

Patch to review:
diff --git a/dask/config.py b/dask/config.py
index 24ed1e3a673..67be157ed7d 100644
--- a/dask/config.py
+++ b/dask/config.py
@@ -9,6 +9,8 @@
 except ImportError:
     yaml = None
 
+from collections import Mapping
+
 from .compatibility import makedirs
 
 
@@ -434,4 +436,26 @@ def update_defaults(new, config=config, defaults=defaults):
     update(config, new, priority='old')
 
 
+def expand_environment_variables(config):
+    ''' Utility to expand environment variables in a config dictionary
+
+    This function will recursively search through any nested dictionaries
+    and/or lists.
+
+    config : dict, iterable, or str
+        Input object to search for environment variables
+    '''
+    if isinstance(config, Mapping):
+        for k, v in config.items():
+            config[k] = expand_environment_variables(v)
+    elif isinstance(config, str):
+        config = os.path.expandvars(config)
+    elif isinstance(config, (list, tuple, set)):
+        config = [expand_environment_variables(v) for v in config]
+    else:
+        # fail gracefully for objects that don't fall into these catagories
+        pass
+    return config
+
+
 refresh()
diff --git a/dask/tests/test_config.py b/dask/tests/test_config.py
index cc47e5a5181..d09538a6cc8 100644
--- a/dask/tests/test_config.py
+++ b/dask/tests/test_config.py
@@ -6,7 +6,7 @@
 import dask.config
 from dask.config import (update, merge, collect, collect_yaml, collect_env,
                          get, ensure_file, set, config, rename,
-                         update_defaults, refresh)
+                         update_defaults, refresh, expand_environment_variables)
 from dask.utils import tmpfile
 
 
@@ -279,3 +279,18 @@ def test_refresh():
 
     refresh(paths=[], env={'DASK_C': '3'}, config=config, defaults=defaults)
     assert config == {'a': 1, 'c': 3}
+
+
+def test_expand_environment_variables():
+
+    os.environ["ENV_A"] = "foo"
+    config = '$ENV_A'
+    assert expand_environment_variables(config) == 'foo'
+
+    os.environ["ENV_B"] = "bar"
+    config = {'a': '$ENV_B'}
+    assert expand_environment_variables(config)['a'] == 'bar'
+
+    os.environ["ENV_C"] = "spam"
+    config = {'a': 'A', 'b': ['1', 2, '$ENV_C']}
+    assert expand_environment_variables(config)['b'] == ['1', 2, 'spam']


Code review comments:
diff_hunk: @@ -434,4 +436,26 @@ def update_defaults(new, config=config, defaults=defaults):
     update(config, new, priority='old')
 
 
+def expand_environment_variables(config):
+    ''' Utility to expand environment variables in a config dictionary
+
+    This function will recursively search through any nested dictionaries
+    and/or lists.
+
+    config : dict, iterable, or str
+        Input object to search for environment variables
+    '''
+    if isinstance(config, Mapping):
+        for k, v in config.items():
+            config[k] = expand_environment_variables(v)
+    elif isinstance(config, str):
+        config = os.path.expandvars(config)
+    elif isinstance(config, (list, tuple, set)):
+        config = [expand_environment_variables(v) for v in config]
path: dask/config.py
text: @user1:
We might call `type(config)(...)` on the result to maintain type

diff_hunk: @@ -279,3 +279,18 @@ def test_refresh():
 
     refresh(paths=[], env={'DASK_C': '3'}, config=config, defaults=defaults)
     assert config == {'a': 1, 'c': 3}
+
+
+def test_expand_environment_variables():
+
+    os.environ["ENV_A"] = "foo"
+    config = '$ENV_A'
+    assert expand_environment_variables(config) == 'foo'
+
+    os.environ["ENV_B"] = "bar"
+    config = {'a': '$ENV_B'}
+    assert expand_environment_variables(config)['a'] == 'bar'
+
+    os.environ["ENV_C"] = "spam"
+    config = {'a': 'A', 'b': ['1', 2, '$ENV_C']}
+    assert expand_environment_variables(config)['b'] == ['1', 2, 'spam']
path: dask/tests/test_config.py
text: @user1:
We should probably clean up `os.environ` after running this test in a try-finally block.  Also, we can probably use `ENV_A` in all cases.