diff --git a/alakazam.py b/alakazam.py index 204d2c5..53a8151 100755 --- a/alakazam.py +++ b/alakazam.py @@ -10,7 +10,9 @@ from tabulate import tabulate import click import dotenv from jinja2 import Environment, FileSystemLoader -import yaml +from ruamel.yaml import YAML +from ruamel.yaml.constructor import SafeConstructor +from ruamel.yaml.nodes import ScalarNode import re COMBINE_PATH = os.path.dirname(os.path.realpath(__file__)) + "/combine.yml" @@ -18,13 +20,30 @@ CONFIG_FILE_NAME = 'alaka.yaml' CONFIGS = {} +"""Preserve quotes while parsing yaml files""" +class MySafeConstructor(SafeConstructor): + def construct_yaml_str(self, node): + value = self.construct_scalar(node) + if isinstance(node, ScalarNode) and (node.style == '"' or node.style == "'"): + # If node was quoted, keep quotes in the returned value + return '{}{}{}'.format(node.style, value, node.style) + # Otherwise, return value as is + return value + +MySafeConstructor.add_constructor( + 'tag:yaml.org,2002:str', + MySafeConstructor.construct_yaml_str) + + def read_config(filepath): filepath = Path(filepath).expanduser() if not filepath.exists(): logging.warning(f"config file {filepath} does not exist") return {} + yaml = YAML(typ='safe', pure=True) # Set type to 'safe' and use pure Python mode + yaml.Constructor = MySafeConstructor with open(filepath) as file: - yaml_config = yaml.safe_load(file) + yaml_config = yaml.load(file) if not yaml_config: logging.warning(f"config file {filepath} is empty") return {} @@ -32,7 +51,7 @@ def read_config(filepath): jinja = Environment(loader=FileSystemLoader( ['/', '.']), trim_blocks=True, lstrip_blocks=True) template = jinja.get_template(filepath.as_posix(), globals=globals) - return yaml.safe_load(template.render()) + return yaml.load(template.render()) """ Get value from nested dicts, return None if one of the keys does not exists """