From 06fc4fe58789f872929cde8587c30aa272f24ed9 Mon Sep 17 00:00:00 2001
From: mikecooke77 <michael.c.cooke@metoffice.gov.uk>
Date: Fri, 27 Sep 2024 09:57:51 +0100
Subject: [PATCH] Update the method for variable substiution to include
 environment variables

---
 docs/cli.rst                                  |  8 +++-
 src/yamlprocessor/datapreprocessor.py         | 37 +++++++++++++++----
 .../tests/test_datapreprocess.py              | 26 ++++++++++---
 3 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/docs/cli.rst b/docs/cli.rst
index 14047a0..fa469c1 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -92,7 +92,13 @@ See :doc:`data-preprocessor` for detail.
 
 .. option:: --define=KEY=VALUE, -D KEY=VALUE
 
-   Map KEY to VALUE for variable substitutions.
+   Map KEY to VALUE for variable substitutions. These override
+   environment variables which are used by default in the variable
+   substitution.
+
+.. option:: --no-environment, -i
+
+   Do not use environment variables in variable substitutions.
 
 yp-schema
 ---------
diff --git a/src/yamlprocessor/datapreprocessor.py b/src/yamlprocessor/datapreprocessor.py
index f1a3dc0..0de05af 100755
--- a/src/yamlprocessor/datapreprocessor.py
+++ b/src/yamlprocessor/datapreprocessor.py
@@ -12,16 +12,30 @@
 """
 
 import argparse
+import os
+import re
 import sys
 
 
 class DataPreProcessor:
 
     def __init__(self):
-        self.replacements = {}
+        self.replacements = os.environ.copy()
+
+    def __replace_placeholders(self, text):
+        # Create a regex pattern that matches $VAR or ${VAR}
+        pattern = re.compile(r'\$\{(\w+)\}|\$(\w+)')
+
+        # Function to get the replacement value from env_vars
+        def replacer(match):
+            var_name = match.group(1) or match.group(2)
+            return self.replacements.get(var_name, match.group(0))
+
+        # Substitute the placeholders with actual values
+        return pattern.sub(replacer, text)
 
     def add_replacements_map(self, replacements):
-        self.replacements = replacements
+        self.replacements.update(replacements)
 
     def process_yaml(self, in_yaml, out_yaml):
         # read yaml file
@@ -37,9 +51,8 @@ def process_yaml(self, in_yaml, out_yaml):
                 # retrieve header file
                 yaml_header_File = iline.split('=')[1].rstrip()
                 # replace variables in the string
-                for key, value in self.replacements.items():
-                    yaml_header_File = \
-                        yaml_header_File.replace(f'${key}', value)
+                yaml_header_File = self.__replace_placeholders(
+                    yaml_header_File)
                 # open header file
                 with open(yaml_header_File, 'r') as file:
                     auxFileData = file.read()
@@ -72,12 +85,17 @@ def main():
         help='Name of output file, "-" for STDOUT'
     )
 
-    # Optional --define arguments
+    # Optional
     parser.add_argument(
-        '--define', '-d',
+        '--define', '-D',
         action='append',
         help='Key-value pairs in the format key=value', default=[]
     )
+    parser.add_argument(
+        '--no-environment', '-i',
+        action='store_true',
+        default=False,
+        help='Do not use environment variables in variable substitutions')
 
     # Parse arguments and print for sanity checking
     args = parser.parse_args()
@@ -85,7 +103,8 @@ def main():
     print(f"Output file: {args.output_file}", file=sys.stderr)
     print(f"Defines: {args.define}", file=sys.stderr)
 
-    # Process define arguments into a dictionary for passing to the class
+    # Process define arguments into a dictionary for adding to the
+    # environment variable dictionary
     key_value_pairs = {}
     if args.define:
         for item in args.define:
@@ -94,6 +113,8 @@ def main():
 
     # Run preprocessor
     preprocessor = DataPreProcessor()
+    if args.no_environment:
+        preprocessor.replacements.clear()
     preprocessor.add_replacements_map(key_value_pairs)
     preprocessor.process_yaml(args.input_file, args.output_file)
 
diff --git a/src/yamlprocessor/tests/test_datapreprocess.py b/src/yamlprocessor/tests/test_datapreprocess.py
index 838ada4..9bbd6c5 100644
--- a/src/yamlprocessor/tests/test_datapreprocess.py
+++ b/src/yamlprocessor/tests/test_datapreprocess.py
@@ -20,6 +20,13 @@ def test_main_0(tmp_path, yaml):
     tel: *groups
 """
     yaml_1 = """
+DIRECT_INCLUDE=${FILE_PATH}/aux.yaml
+
+data:
+    brain: *banana
+    tel: *groups
+"""
+    yaml_2 = """
 _:
 - &banana 1
 - &groups [4, 5, 6]
@@ -37,16 +44,25 @@ def test_main_0(tmp_path, yaml):
     with infilename.open('w') as infile:
         infile.write(yaml_0)
 
+    infilename = tmp_path / 'in_1.yaml'
+    with infilename.open('w') as infile:
+        infile.write(yaml_1)
+
     auxfilename = tmp_path / 'aux.yaml'
     with auxfilename.open('w') as auxfile:
-        auxfile.write(yaml_1)
+        auxfile.write(yaml_2)
 
     # Run preprocessor
     preprocessor = DataPreProcessor()
     keymap = {"FILE_PATH": str(tmp_path)}
     preprocessor.add_replacements_map(keymap)
-    outfilename = tmp_path / 'test_0.yaml'
-    preprocessor.process_yaml(tmp_path / 'in_0.yaml', outfilename)
-    # Check output
+    # Setup reference
     ref_yaml = yaml.load(reference)
-    assert yaml.load(outfilename.open()) == ref_yaml
+    # Test first style input
+    outfilename0 = tmp_path / 'test_0.yaml'
+    preprocessor.process_yaml(tmp_path / 'in_0.yaml', outfilename0)
+    assert yaml.load(outfilename0.open()) == ref_yaml
+    # Test second style input
+    outfilename1 = tmp_path / 'test_1.yaml'
+    preprocessor.process_yaml(tmp_path / 'in_1.yaml', outfilename1)
+    assert yaml.load(outfilename1.open()) == ref_yaml