import yaml

from parser.parser_base import ParserBase


class YAMLParser(ParserBase):
    def parse(self):
        return self._parse_yaml()

    def _parse_yaml(self):
        with open(self.filepath, "r") as file:
            dag_config_dict = yaml.safe_load(file)

        dag = self._try_get(dag_config_dict, "dag")

        initial_parser = {
            "dag_id": self._try_get(dag, "id"),
            "data_path": self._try_get(dag, "data_path"),
            "output_folder": self._try_get(dag, "output_folder"),
            "description": self._try_get(dag, "description"),
            "feature_engineering": self._try_get(dag, "feature_engineering"),
            "model": self._try_get(dag, "model"),
        }

        return initial_parser
