diff --git a/samcli/lib/config/file_manager.py b/samcli/lib/config/file_manager.py index 146946056f..49a60fdadd 100644 --- a/samcli/lib/config/file_manager.py +++ b/samcli/lib/config/file_manager.py @@ -9,6 +9,8 @@ from typing import Any import tomlkit +from ruamel.yaml import YAML, YAMLError +from ruamel.yaml.compat import StringIO from samcli.lib.config.exceptions import FileParseException @@ -159,3 +161,102 @@ def put_comment(document: dict, comment: str) -> Any: def _to_toml(document: dict) -> tomlkit.TOMLDocument: """Ensure that a dictionary-like object is a TOMLDocument.""" return tomlkit.parse(tomlkit.dumps(document)) + + +class YamlFileManager(FileManager): + """ + Static class to read and write yaml files. + """ + + yaml = YAML() + file_format = "YAML" + + @staticmethod + def read(filepath: Path) -> Any: + """ + Read a YAML file at the given path. + + Parameters + ---------- + filepath: Path + The Path object that points to the file to be read. + + Returns + ------- + Any + A dictionary-like yaml object, which represents the contents of the YAML file at the + provided location. + """ + yaml_doc = YamlFileManager.yaml.load("") + try: + yaml_doc = YamlFileManager.yaml.load(filepath.read_text()) + except OSError as e: + LOG.debug(f"OSError occurred while reading {YamlFileManager.file_format} file: {str(e)}") + except YAMLError as e: + raise FileParseException(e) from e + + return yaml_doc + + @staticmethod + def write(document: dict, filepath: Path): + """ + Write the contents of a dictionary to a YAML file at the provided location. + + Parameters + ---------- + document: dict + The object to write. + filepath: Path + The final location for the YAML file to be written. + """ + if not document: + LOG.debug("No document given to YamlFileManager to write.") + return + + yaml_doc = YamlFileManager._to_yaml(document) + + if yaml_doc.get(COMMENT_KEY, None): # Comment appears at the top of doc + yaml_doc.yaml_set_start_comment(document[COMMENT_KEY]) + yaml_doc.pop(COMMENT_KEY) + + YamlFileManager.yaml.dump(yaml_doc, filepath) + + @staticmethod + def put_comment(document: Any, comment: str) -> Any: + """ + Put a comment in a document object. + + Parameters + ---------- + document: Any + The yaml object to write + comment: str + The comment to include in the document. + + Returns + ------- + Any + The new yaml document, with the comment added to it. + """ + document = YamlFileManager._to_yaml(document) + document.yaml_set_start_comment(comment) + return document + + @staticmethod + def _to_yaml(document: dict) -> Any: + """ + Ensure a dictionary-like object is a YAML document. + + Parameters + ---------- + document: dict + A dictionary-like object to parse. + + Returns + ------- + Any + A dictionary-like YAML object, as derived from `yaml.load()`. + """ + with StringIO() as stream: + YamlFileManager.yaml.dump(document, stream) + return YamlFileManager.yaml.load(stream.getvalue()) diff --git a/samcli/lib/config/samconfig.py b/samcli/lib/config/samconfig.py index d70b8fe1bb..eb8f0f8f1c 100644 --- a/samcli/lib/config/samconfig.py +++ b/samcli/lib/config/samconfig.py @@ -8,7 +8,7 @@ from typing import Any, Dict, Iterable, Type from samcli.lib.config.exceptions import FileParseException, SamConfigFileReadException, SamConfigVersionException -from samcli.lib.config.file_manager import FileManager, TomlFileManager +from samcli.lib.config.file_manager import FileManager, TomlFileManager, YamlFileManager from samcli.lib.config.version import SAM_CONFIG_VERSION, VERSION_KEY LOG = logging.getLogger(__name__) @@ -26,6 +26,8 @@ class SamConfig: FILE_MANAGER_MAPPER: Dict[str, Type[FileManager]] = { ".toml": TomlFileManager, + ".yaml": YamlFileManager, + ".yml": YamlFileManager, } def __init__(self, config_dir, filename=None): diff --git a/tests/unit/commands/samconfig/test_samconfig.py b/tests/unit/commands/samconfig/test_samconfig.py index c9d7e442fc..b71d16a1c1 100644 --- a/tests/unit/commands/samconfig/test_samconfig.py +++ b/tests/unit/commands/samconfig/test_samconfig.py @@ -14,7 +14,9 @@ from unittest import TestCase from unittest.mock import patch, ANY import logging +from parameterized import parameterized from samcli.lib.config.exceptions import SamConfigFileReadException +from samcli.lib.config.file_manager import YamlFileManager from samcli.lib.config.samconfig import SamConfig, DEFAULT_ENV, TomlFileManager from samcli.lib.utils.packagetype import ZIP, IMAGE @@ -1254,13 +1256,20 @@ def test_file_manager_not_declared(self): with self.assertRaises(SamConfigFileReadException): SamConfig(config_path, filename="samconfig") - def test_file_manager_toml(self): + @parameterized.expand( + [ + ("samconfig.toml", TomlFileManager), + ("samconfig.yaml", YamlFileManager), + ("samconfig.yml", YamlFileManager), + ] + ) + def test_file_manager(self, filename, expected_file_manager): config_dir = tempfile.gettempdir() - config_path = Path(config_dir, "samconfig.toml") + config_path = Path(config_dir, filename) - samconfig = SamConfig(config_path, filename="samconfig.toml") + samconfig = SamConfig(config_path, filename=filename) - self.assertIs(samconfig.file_manager, TomlFileManager) + self.assertIs(samconfig.file_manager, expected_file_manager) @contextmanager diff --git a/tests/unit/lib/samconfig/test_file_manager.py b/tests/unit/lib/samconfig/test_file_manager.py index 0973637023..d6b4fe82c0 100644 --- a/tests/unit/lib/samconfig/test_file_manager.py +++ b/tests/unit/lib/samconfig/test_file_manager.py @@ -3,9 +3,10 @@ from unittest import TestCase import tomlkit -from samcli.lib.config.exceptions import FileParseException +from ruamel.yaml import YAML -from samcli.lib.config.file_manager import COMMENT_KEY, TomlFileManager +from samcli.lib.config.exceptions import FileParseException +from samcli.lib.config.file_manager import COMMENT_KEY, TomlFileManager, YamlFileManager class TestTomlFileManager(TestCase): @@ -48,6 +49,7 @@ def test_write_toml(self): self.assertIn("[config_env.topic2.parameters]", txt) self.assertIn('word = "clarity"', txt) self.assertIn("# This is a comment", txt) + self.assertNotIn(COMMENT_KEY, txt) def test_dont_write_toml_if_empty(self): config_dir = tempfile.gettempdir() @@ -99,3 +101,81 @@ def test_toml_put_comment(self): txt = tomlkit.dumps(toml_doc) self.assertIn("# This is a comment", txt) + + +class TestYamlFileManager(TestCase): + + yaml = YAML() + + def test_read_yaml(self): + config_dir = tempfile.gettempdir() + config_path = Path(config_dir, "samconfig.yaml") + config_path.write_text("version: 0.1\nconfig_env:\n topic1:\n parameters:\n word: clarity\n") + + config_doc = YamlFileManager.read(config_path) + + self.assertEqual( + config_doc, + {"version": 0.1, "config_env": {"topic1": {"parameters": {"word": "clarity"}}}}, + ) + + def test_read_yaml_invalid_yaml(self): + config_dir = tempfile.gettempdir() + config_path = Path(config_dir, "samconfig.yaml") + config_path.write_text("fake: not real\nthisYaml isn't correct") + + with self.assertRaises(FileParseException): + YamlFileManager.read(config_path) + + def test_read_yaml_file_path_not_valid(self): + config_dir = "path/that/doesnt/exist" + config_path = Path(config_dir, "samconfig.yaml") + + config_doc = YamlFileManager.read(config_path) + + self.assertEqual(config_doc, self.yaml.load("")) + + def test_write_yaml(self): + config_dir = tempfile.gettempdir() + config_path = Path(config_dir, "samconfig.yaml") + yaml = { + "version": 0.1, + "config_env": {"topic2": {"parameters": {"word": "clarity"}}}, + COMMENT_KEY: "This is a comment", + } + + YamlFileManager.write(yaml, config_path) + + txt = config_path.read_text() + self.assertIn("version: 0.1", txt) + self.assertIn("config_env:\n topic2:\n parameters:\n", txt) + self.assertIn("word: clarity", txt) + self.assertIn("# This is a comment", txt) + self.assertNotIn(COMMENT_KEY, txt) + + def test_dont_write_yaml_if_empty(self): + config_dir = tempfile.gettempdir() + config_path = Path(config_dir, "samconfig.yaml") + config_path.write_text("nothing to see here\n") + yaml = {} + + YamlFileManager.write(yaml, config_path) + + self.assertEqual(config_path.read_text(), "nothing to see here\n") + + def test_write_yaml_file_bad_path(self): + config_path = Path("path/to/some", "file_that_doesnt_exist.yaml") + + with self.assertRaises(FileNotFoundError): + YamlFileManager.write(self.yaml.load("key: some value"), config_path) + + def test_yaml_put_comment(self): + config_dir = tempfile.gettempdir() + config_path = Path(config_dir, "samconfig.yaml") + yaml_doc = self.yaml.load("version: 0.1\nconfig_env:\n topic2:\n parameters:\n word: clarity\n") + + yaml_doc = YamlFileManager.put_comment(yaml_doc, "This is a comment") + + self.yaml.dump(yaml_doc, config_path) + txt = config_path.read_text() + self.assertIn("# This is a comment", txt)