Repository: isi-vista/vistautils

Issue:
Allow parameters to be interpolated from environmental variables

Comments:

Patch:
diff --git a/vistautils/parameters.py b/vistautils/parameters.py
index 3637f60..71bbfdf 100644
--- a/vistautils/parameters.py
+++ b/vistautils/parameters.py
@@ -1,5 +1,6 @@
 import inspect
 import logging
+import os
 import re
 import shutil
 from pathlib import Path
@@ -93,6 +94,12 @@ class Parameters:
                 ret.put(key, val)
         return Parameters(ret.build())
 
+    def as_mapping(self) -> Mapping[str, Any]:
+        """
+        Get these parameter values as a ``Mapping``.
+        """
+        return self._data
+
     def creatable_directory(self, param: str) -> Path:
         """
         Get a directory which can be written to.
@@ -678,6 +685,7 @@ class Parameters:
         return current
 
 
+@attrs(auto_attribs=True)
 class YAMLParametersLoader:
     """
     Loads `Parameters` from a modified YAML format.
@@ -698,53 +706,138 @@ class YAMLParametersLoader:
     string (and the `%`s) with the value of the parameter.  Failed interpolations will result
     in a `ParameterError`.
 
+    If *interpolate_environmental_variables* (default: false) is specified, then environmental
+    variables will be available for interpolation, though they will not themselves appear in the
+    loaded parameters.  Explicitly-specified parameters have priority over environmental variables.
+
     See unit tests in `test_parameters` for examples.
     """
 
+    interpolate_environmental_variables: bool = False
+
     def load(
-        self, f: Union[str, Path], context=Parameters.empty(), root_path: Path = None
+        self,
+        f: Union[str, Path],
+        context=None,
+        root_path: Path = None,
+        *,
+        included_context: Parameters = None,
     ):
         """
         Loads parameters from a YAML file.
 
-        If `context` is specified, its content will be included in the returned Parameters (if
-        not overridden) and will be available for interpolation.
+        If `included_context` is specified, its content will be included in the returned
+        Parameters (if not overridden) and will be available for interpolation.
 
         If `root_path` is specified, that path is used for resolving relative path names in
         includes instead of the path of the parameter file being loaded.
         """
+
+        # handle deprecated context parameter
+        if context is not None:
+            if included_context is not None:
+                raise ParameterError(
+                    "Cannot specify both included_context and deprecated context"
+                    "parameters. Only specify the former."
+                )
+            else:
+                included_context = context
+
         if isinstance(f, str):
             f = Path(f)
-        try:
-            # if no special path is specified, included files will be resolved relative to
-            # this file's path
-            if not root_path:
-                root_path = f.parent
 
-            with open(f, "r") as ymlfile:
-                raw_yaml = yaml.safe_load(ymlfile)
-                self._validate(raw_yaml)
-            cur_context = context
+        # if no special path is specified, included files will be resolved relative to
+        # this file's path
+        if not root_path:
+            root_path = f.parent
+
+        return self._inner_load(
+            f.read_text(encoding="utf-8"),
+            error_string=str(f),
+            root_path=root_path,
+            included_context=included_context,
+        )
+
+    def load_string(
+        self,
+        param_file_content: str,
+        *,
+        included_context: Parameters = Parameters.empty(),
+    ) -> Parameters:
+        """
+        Loads parameters from a string.
+
+        This behaves just like *load*, except relative includes are not allowed.
+        """
+        return self._inner_load(
+            param_file_content,
+            error_string=f"String param file:\n{param_file_content}",
+            root_path=None,
+            included_context=included_context,
+        )
+
+    def _inner_load(
+        self,
+        param_file_content: str,
+        error_string: str,
+        *,
+        included_context=Parameters.empty(),
+        root_path: Optional[Path] = None,
+    ):
+        """
+        Loads parameters from a YAML file.
+
+        If `context` is specified, its content will be included in the returned Parameters (if
+        not overridden) and will be available for interpolation.
+
+        If `root_path` is specified, that path is used for resolving relative path names in
+        includes instead of the path of the parameter file being loaded.
+        """
+        try:
+            raw_yaml = yaml.safe_load(param_file_content)
+            self._validate(raw_yaml)
+            previously_loaded = included_context
 
             # process and remove special include directives
             if "_includes" in raw_yaml:
                 for included_file in raw_yaml["_includes"]:
                     _logger.info("Processing included parameter file %s", included_file)
-                    included_file_path = Path(root_path, included_file).resolve()
-                    cur_context = self._unify(
-                        cur_context,
+                    if not os.path.isabs(included_file):
+                        if root_path is not None:
+                            included_file_path = Path(root_path, included_file).resolve()
+                        else:
+                            raise ParameterError(
+                                "Cannot do relative includes when loading from"
+                                "a string."
+                            )
+                    else:
+                        included_file_path = Path(included_file)
+                    previously_loaded = self._unify(
+                        previously_loaded,
                         self.load(
-                            included_file_path, root_path=root_path, context=cur_context
+                            included_file_path,
+                            root_path=root_path,
+                            context=previously_loaded,
                         ),
                     )
                 del raw_yaml["_includes"]
 
+            interpolation_context = dict(included_context.as_mapping())
+            if self.interpolate_environmental_variables:
+                for k, v in os.environ.items():
+                    # environmental variables are overridden by explicit parameters
+                    if k not in interpolation_context:
+                        interpolation_context[k] = v
+
             return self._unify(
-                cur_context,
-                self._interpolate(Parameters.from_mapping(raw_yaml), cur_context),
+                previously_loaded,
+                self._interpolate(
+                    Parameters.from_mapping(raw_yaml),
+                    Parameters.from_mapping(interpolation_context),
+                ),
             )
         except Exception as e:
-            raise IOError("Failure while loading parameter file " + str(f)) from e
+            raise IOError(f"Failure while loading parameter file {error_string}") from e
 
     @staticmethod
     def _validate(raw_yaml: Mapping):
@@ -773,6 +866,7 @@ class YAMLParametersLoader:
 
     _INTERPOLATION_REGEX = re.compile(r"%([\w\.]+)%")
 
+    # noinspection PyProtectedMember
     @staticmethod
     def _interpolate(to_interpolate: Parameters, context: Parameters) -> Parameters:
         """Perform interpolation within arbitrarily nested Parameters, looking up values
@@ -826,7 +920,9 @@ class YAMLParametersLoader:
                 replaced_end_node = end_node
             interpolation_mapping[start_node] = replaced_end_node
         return Parameters.from_mapping(
-            {key: interpolation_mapping[key] for key in to_interpolate._data}
+            immutabledict(
+                (key, interpolation_mapping[key]) for key in to_interpolate._data.keys()
+            )
         )
 
     def _unify(self, old: Parameters, new: Parameters, namespace="") -> Parameters:
diff --git a/vistautils/range.py b/vistautils/range.py
index b80fb06..7b9bcc7 100644
--- a/vistautils/range.py
+++ b/vistautils/range.py
@@ -948,7 +948,7 @@ class _SortedDictRangeSet(RangeSet[T], metaclass=ABCMeta):
         # If we would insert at the end (are greater than all the elements, the only range that
         # could possibly overlap is the last one.
         if from_index == len(rlb):
-            last_range: Range[T] = rlb[rlb.iloc[-1]]
+            last_range: Range[T] = rlb[rlb.keys()[-1]]
             if last_range.intersects(rng):
                 return immutableset([last_range])
             return immutableset()
@@ -956,17 +956,17 @@ class _SortedDictRangeSet(RangeSet[T], metaclass=ABCMeta):
         # If we would insert at the start (are smaller than all the elements, the only range that
         # could possibly overlap is the first one.
         if to_index == 0:
-            first_range: Range[T] = rlb[rlb.iloc[0]]
+            first_range: Range[T] = rlb[rlb.keys()[0]]
             if first_range.intersects(rng):
                 return immutableset([first_range])
             return immutableset()
         return immutableset(
             [
-                rlb[rlb.iloc[index]]
+                rlb[rlb.keys()[index]]
                 # The ranges at the extreme indices do not necessarily overlap,
                 for index in range(max(0, from_index - 1), min(to_index, len(rlb) - 1))
                 # so this explicit check is needed.
-                if rlb[rlb.iloc[index]].intersects(rng)
+                if rlb[rlb.keys()[index]].intersects(rng)
             ]
         )
 
@@ -988,13 +988,13 @@ class _SortedDictRangeSet(RangeSet[T], metaclass=ABCMeta):
         if containing_or_below_index >= 0:
             # if such a set exists, we need to check if we are contained in it...
             latest_beginning_before = sorted_dict[
-                sorted_dict.iloc[containing_or_below_index]
+                sorted_dict.keys()[containing_or_below_index]
             ]
             if limit_as_bound <= latest_beginning_before._upper_bound:
                 return latest_beginning_before
 
         if idx < len(sorted_dict):
-            return sorted_dict[sorted_dict.iloc[idx]]
+            return sorted_dict[sorted_dict.keys()[idx]]
         else:
             return None
 
@@ -1301,7 +1301,7 @@ def _value_below(sorted_dict: SortedDict, key: T) -> Optional[Any]:
     if idx >= 0:
         if idx >= len(sorted_dict):
             idx = len(sorted_dict) - 1
-        lb_key = sorted_dict.iloc[idx]
+        lb_key = sorted_dict.keys()[idx]
         return sorted_dict[lb_key]
     else:
         return None
@@ -1318,9 +1318,9 @@ def _value_at_or_below(sorted_dict: SortedDict, key: T) -> Optional[Any]:
 
     idx = sorted_dict.bisect_left(key)
 
-    if idx >= len(sorted_dict) or key != sorted_dict.iloc[idx]:
+    if idx >= len(sorted_dict) or key != sorted_dict.keys()[idx]:
         if idx > 0:
-            key = sorted_dict.iloc[idx - 1]
+            key = sorted_dict.keys()[idx - 1]
         else:
             return None
     return sorted_dict[key]
@@ -1333,7 +1333,7 @@ def _value_at_or_above(sorted_dict: SortedDict, key: T) -> Optional[Any]:
     if idx >= len(sorted_dict):
         return None
 
-    return sorted_dict[sorted_dict.iloc[idx]]
+    return sorted_dict[sorted_dict.keys()[idx]]
 
 
 def _clear(


Test patch:
diff --git a/tests/test_parameters.py b/tests/test_parameters.py
index 02ab6e5..f19ea3d 100644
--- a/tests/test_parameters.py
+++ b/tests/test_parameters.py
@@ -2,6 +2,7 @@ import os
 import shutil
 import tempfile
 import yaml
+from textwrap import dedent
 from pathlib import Path
 from unittest import TestCase
 
@@ -18,15 +19,19 @@ from vistautils._graph import ParameterInterpolationError
 
 
 class TestParameters(TestCase):
-    WRITING_REFERENCE = """hello: world
-moo:
-    nested_dict:
-        lalala: fooo
-        list:
-        - 1
-        - 2
-        - 3
-        meep: 2\n"""
+    WRITING_REFERENCE = dedent(
+        """\
+            hello: world
+            moo:
+                nested_dict:
+                    lalala: fooo
+                    list:
+                    - 1
+                    - 2
+                    - 3
+                    meep: 2
+        """
+    )
 
     def test_writing_to_yaml(self):
         params = Parameters.from_mapping(
@@ -132,20 +137,23 @@ moo:
         ):
             params.floating_point("test_float", valid_range=Range.open(0.0, 1.0))
 
-    MULTIPLE_INTERPOLATION_REFERENCE = """the_ultimate_fruit: \"%apple%\"
-apple: \"%banana%\"
-banana: \"%pear%\"
-pear: raspberry
-"""
-    MULTIPLE_INTERPOLATION_REFERENCE_NEEDING_CONTEXT = """the_ultimate_fruit: \"%apple%\"
-apple: \"%banana%\"
-banana: \"%pear%\"
-pear: \"raspberry/%hello%\"
-"""
-    NESTED_INTERPOLATION = """key: \"%moo.nested_dict.meep%\"
-key2: \"%moo.nested_dict.lalala%\"
-key3: \"%moo.nested_dict%\"
-"""
+    MULTIPLE_INTERPOLATION_REFERENCE = """
+            the_ultimate_fruit: "%apple%"
+            apple: "%banana%"
+            banana: "%pear%"
+            pear: raspberry
+            """
+    MULTIPLE_INTERPOLATION_REFERENCE_NEEDING_CONTEXT = """
+            the_ultimate_fruit: "%apple%"
+            apple: "%banana%"
+            banana: "%pear%"
+            pear: "raspberry/%hello%"
+            """
+    NESTED_INTERPOLATION = """
+            key: "%moo.nested_dict.meep%"
+            key2: "%moo.nested_dict.lalala%"
+            key3: "%moo.nested_dict%"
+            """
 
     def test_interpolation(self):
         context = Parameters.from_mapping(yaml.safe_load(self.WRITING_REFERENCE))
@@ -204,9 +212,50 @@ key3: \"%moo.nested_dict%\"
 
         with self.assertRaisesRegex(
             ParameterInterpolationError,
-            r"These interpolated parameters form at least one graph cycle that must be fixed: \('b', 'c'\)",
+            r"These interpolated parameters form at least one graph cycle that must be fixed: "
+            r"\('b', 'c'\)",
         ):
             loader._interpolate(
                 Parameters.from_mapping(yaml.safe_load('a: "%b%"\nb: "%c%"\nc: "%b%"')),
                 context,
             )
+
+    def test_environmental_variable_interpolation(self):
+        loader = YAMLParametersLoader(interpolate_environmental_variables=True)
+        os.environ["___TEST_PARAMETERS___"] = "foo"
+        os.environ["___TEST_CLASHING_PARAM___"] = "bar"
+        loaded_params = loader.load_string(ENV_VAR_INTERPOLATION_INPUT)
+
+        reference_params = Parameters.from_mapping(
+            yaml.safe_load(ENV_VAR_INTERPOLATION_REFERENCE)
+        )
+
+        self.assertEqual(reference_params, loaded_params)
+
+    def test_double_context_fail(self):
+        with self.assertRaises(ParameterError):
+            YAMLParametersLoader().load(f=None, context="bar", included_context="baz")
+
+
+# Used by test_environmental_variable_interpolation.
+# Here we test:
+# (a) one uninterpolated parameter
+# (b) one normally interpolated parameter
+# (c) one parameter interpolated with an environmental variable
+# (d) one parameter interpolated with a key which is both explicitly specified and and an
+#         environmental variable, demonstrating that the explicit parameter "wins"
+ENV_VAR_INTERPOLATION_INPUT = """
+        key1: "fred"
+        regular_interpolation: "moo %key1%"
+        env_var_interpolation: "moo %___TEST_PARAMETERS___%"
+        ___TEST_CLASHING_PARAM___: "rab"
+        interpolation_of_clashing_param: "moo %___TEST_CLASHING_PARAM___%"
+        """
+
+ENV_VAR_INTERPOLATION_REFERENCE = """
+        key1: "fred"
+        regular_interpolation: "moo fred"
+        env_var_interpolation: "moo foo"
+        ___TEST_CLASHING_PARAM___: "rab"
+        interpolation_of_clashing_param: "moo rab"
+        """