From d6a06530b662311a04500a57366997bcf643662b Mon Sep 17 00:00:00 2001 From: Carl Benson Date: Wed, 29 May 2024 19:25:20 -0400 Subject: [PATCH] Add functionality for loading both YAML and JSON source files --- scripts/lib/hcl_generator.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/scripts/lib/hcl_generator.py b/scripts/lib/hcl_generator.py index ff107d6..6f94257 100644 --- a/scripts/lib/hcl_generator.py +++ b/scripts/lib/hcl_generator.py @@ -41,10 +41,16 @@ def generate_locals(rules: List[AwsConfigRule], output_file: Path) -> None: format_hcl() -def load_source_file(file_name: str) -> dict: +def load_source_file(file_name: Union[Path, str]) -> dict: """Return a map of rule definitions.""" + # Load using JSON or YAML based on the file extension. with Path(file_name).open() as f: - data = json.loads(f.read()) + if file_name.suffix in ('.yml', '.yaml',): + data = yaml.safe_load(f) + elif file_name.suffix == '.json': + data = json.loads(f.read()) + else: + raise ValueError(f"Unsupported file extension: {file_name.suffix}") return data def format_hcl() -> None: