Merge pull request #664 from yosoyjay/add-variables-to-config-keys

Add ability to use variables in config dictionary keys

pyazhpc/azconfig.py

@@ -1,3 +1,4 @@
+import copy
 import json
 import os
 import re
@@ -14,14 +15,17 @@ class ConfigFile:
         self.data = {}
         self.regex = re.compile(r'({{([^{}]*)}})')
 
-    def open(self, fname):
+    def open(self, fname, preprocess=True, preprocess_extended=False):
         log.debug("opening "+fname)
         self.file_location = os.path.dirname(fname)
         if self.file_location == "":
             self.file_location = "."
         with open(fname) as f:
             self.data = json.load(f)
-    
+
+        if preprocess:
+            self.data = self.preprocess(extended=preprocess_extended)
+
     def save(self, fname):
         with open(fname, "w") as f:
             json.dump(self.data, f, indent=4)
@@ -36,11 +40,27 @@ class ConfigFile:
                 dest = azutil.get_vm_private_ip(self.read_value("resource_group"), install_from)
         log.debug(f"install_from destination : {dest}")
         return dest
-    
+
     def __evaluate_dict(self, x, extended):
+        # Create deep copy of dict to avoid unintentionally deleting items if variables are used in keys
+        x = copy.deepcopy(x)
         ret = {}
+        updated_keys = []
         for k in x.keys():
-            ret[k] = self.__evaluate(x[k], extended)
+            # Update key names if config variable (i.e. {{variables.key}} ) used
+            if "variables." in k:
+                log.debug(f"expanding key {k}")
+                new_key = self.__evaluate(k, extended)
+                ret[new_key] = self.__evaluate(x[k], extended)
+                updated_keys.append(k)
+            else:
+                ret[k] = self.__evaluate(x[k], extended)
+
+        # Delete keys from dict that were updated
+        # - Not doing so can result in duplicate fields in deploy_*json
+        if updated_keys:
+            for old_key in updated_keys:
+                del x[old_key]
         return ret
 
     def __evaluate_list(self, x, extended):
@@ -77,10 +97,10 @@ class ConfigFile:
         except KeyError:
             log.error("read_keys : "+v+" not in config")
             sys.exit(1)
-        
+
         if type(it) is not dict:
             log.error("read_keys : "+v+" is not a dict")
-        
+
         keys = list(it.keys())
         log.debug("read_keys (exit): keys("+v+")="+",".join(keys))
         return keys
@@ -100,7 +120,7 @@ class ConfigFile:
                     else:
                         log.error("invalid path in config file ({v})")
                 it = it[x]
-            
+
             if type(it) is str:
                 res = self.process_value(it)
             else:
@@ -108,7 +128,7 @@ class ConfigFile:
         except KeyError:
             log.debug(f"using default value ({default})")
             res = default
-        
+
         log.debug("read_value (exit): "+v+"="+str(res))
 
         return res
@@ -118,9 +138,9 @@ class ConfigFile:
 
         def repl(match):
             return str(self.process_value(match.group()[2:-2], extended))
-    
+
         v = self.regex.sub(lambda m: str(self.process_value(m.group()[2:-2], extended)), v)
-        
+
         parts = v.split('.')
         prefix = parts[0]
         if len(parts) == 1:
@@ -186,6 +206,6 @@ class ConfigFile:
                     res = f.read()
             else:
                 res = v
-        
+
         log.debug("process_value (exit): "+str(v)+"="+str(res))
         return res

pyazhpc/azhpc.py

@@ -1,4 +1,5 @@
 import argparse
+import copy
 import datetime
 import json
 import os
@@ -8,7 +9,10 @@ import socket
 import sys
 import textwrap
 import time
-import copy
+
+from cryptography.hazmat.backends import default_backend as crypto_default_backend
+from cryptography.hazmat.primitives import serialization as crypto_serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
 
 import arm
 import azconfig
@@ -16,10 +20,6 @@ import azinstall
 import azlog
 import azutil
 
-from cryptography.hazmat.primitives import serialization as crypto_serialization
-from cryptography.hazmat.primitives.asymmetric import rsa
-from cryptography.hazmat.backends import default_backend as crypto_default_backend
-
 log = azlog.getLogger(__name__)
 
 def do_preprocess(args):
@@ -203,7 +203,7 @@ def do_connect(args):
         sshuser = adminuser
     else:
         sshuser = args.user
-    
+
     sshport = c.read_value("ssh_port", 22)
 
     jumpbox = c.read_value("install_from")
@@ -470,7 +470,7 @@ def _wait_for_deployment(resource_group, deploy_name):
                                         elif isinstance(value, list):
                                             wrapped_print(indent, str(key))
                                             pretty_print(value, indent+4)
-                                        else: 
+                                        else:
                                             wrapped_print(indent, f"{key}: {value}")
                                 else:
                                     wrapped_print(indent, str(d))
@@ -565,8 +565,8 @@ def do_slurm_resume(args):
     # first get the resource name
     resource_names, resource_list = _nodelist_expand(args.nodes)
 
-    # Create a copy of the configuration to use as template                                                                 
-    # for the final deployment configuration                                                                                
+    # Create a copy of the configuration to use as template
+    # for the final deployment configuration
     config = copy.deepcopy(config_orig)
     config["resources"] = {}
 
@@ -716,7 +716,7 @@ def do_run_install(args):
     if os.path.isdir(tmpdir):
         log.debug("removing existing tmp directory")
         shutil.rmtree(tmpdir)
-    
+
     adminuser = config["admin_user"]
     private_key_file = adminuser+"_id_rsa"
     public_key_file = adminuser+"_id_rsa.pub"
@@ -914,7 +914,7 @@ if __name__ == "__main__":
     preprocess_parser.set_defaults(func=do_preprocess)
 
     run_parser = subparsers.add_parser(
-        "run", 
+        "run",
         parents=[gopt_parser],
         add_help=False,
         description="run a command on the specified resources",
@@ -940,7 +940,7 @@ if __name__ == "__main__":
     )
 
     scp_parser = subparsers.add_parser(
-        "scp", 
+        "scp",
         parents=[gopt_parser],
         add_help=False,
         description="secure copy",

pyazhpc/test/test_azconfig.py

@@ -20,5 +20,8 @@ class TestConfigFile(unittest.TestCase):
     def test_replace_double_curly_braces(self):
         self.assertEqual(self.config.read_value("double_curly_braces"), "simple_variable=42")
 
+    def test_replace_in_dict_key(self):
+        self.assertEqual(self.config.read_value("resources.foo"), "bar")
+
 if __name__ == "__main__":
     unittest.main()

pyazhpc/test/test_config_file.json

@@ -4,6 +4,10 @@
     "use_variable": "variables.simple_variable",
     "double_curly_braces": "simple_variable={{variables.simple_variable}}",
     "variables": {
-        "simple_variable": 42
+        "simple_variable": 42,
+        "resource_name": "foo"
+    },
+    "resources": {
+        "variables.resource_name": "bar"
     }
 }