modify to import JSON, and not YAML

Config_builder.py

@@ -1,6 +1,7 @@
 from jinja2 import Environment, FileSystemLoader
 import pprint
 import os
+import json
 
 from yaml import safe_load
 try:
@@ -8,10 +9,22 @@ try:
 except ImportError:
     from yaml import Loader
 
-def main(template_filepath,yaml_filepath,output_filepath):
+def main(template_filepath,input_filepath,output_filepath):
+    
+    with open(input_filepath, 'r') as file:
+        input_devices = file.read()
+    
+    device_info = dict()
+
+    if os.path.splitext(input_filepath)[1] == ".json":
+        device_info = json.loads(input_devices)
+    # elif os.path.splitext(input_filepath) == ".yaml":
+    #     devices = safe_load(input_devices)
+    else:
+        print(os.path.splitext(input_filepath))
+        print('{} is not a valid JSON format!\n'.format(input_filepath))
+        return 406
     
-    with open(yaml_filepath, 'r') as file:
-        devices = safe_load(file)
     
     template_path = os.path.dirname(os.path.abspath(template_filepath))
     template_filename = os.path.basename(template_filepath)
@@ -28,20 +41,23 @@ def main(template_filepath,yaml_filepath,output_filepath):
     #output = open(filename, mode="w", encoding="utf-8")        
     output = open(output_filepath, mode="w", encoding="utf-8")        
 
-    
-    for host, keys in devices.items():
-        device = "{"
-        for key, value in keys.items():
-            device += "'{}':'{}', ".format(key,value)
-        device += "'Hostname':'{}' }}".format(host)
-        #pprint.pprint(device)
+    content = template.render(device_info)
 
-        device_dict = eval(device)
-        #pprint.pprint(device_dict)
-        content = template.render( device_dict )
+    output.write('{}\n'.format(content))
+    
+    # for host, keys in devices.items():
+    #     device = "{"
+    #     for key, value in keys.items():
+    #         device += "'{}':'{}', ".format(key,value)
+    #     device += "'Hostname':'{}' }}".format(host)
+    #     #pprint.pprint(device)
+
+    #     device_dict = eval(device)
+    #     #pprint.pprint(device_dict)
+    #     content = template.render( device_dict )
                 
-        print('{}\n'.format(content))
-        output.write('{}\n'.format(content))
+    #     print('{}\n'.format(content))
+    #     output.write('{}\n'.format(content))
     
     return 0
 
@@ -56,8 +72,8 @@ if __name__ == "__main__":
                         required=True
                         )
 
-    parser.add_argument('-yaml',
-                        help='YAML file with devices and config variables',
+    parser.add_argument('-input',
+                        help='Input file with devices and config variables',
                         required=True
                         )
 
@@ -69,8 +85,8 @@ if __name__ == "__main__":
     custom_args = parser.parse_known_args()[0]
     
     template = custom_args.template
-    yaml = custom_args.yaml
+    input = custom_args.input
     output = custom_args.output
         
-    main(template,yaml,output)
+    main(template,input,output)