Merge pull request #112 from proactiveops/fix-yaml-parser

Fix parsing yaml spec files containing + as value

picofun/spec.py

@@ -100,6 +100,20 @@ def parse(self, content: str) -> dict:
 class YAMLSpecParser(SpecParser):
     """Parses an OpenAPI spec file in YAML format."""
 
+    def patch_loader(self, loader: yaml.SafeLoader) -> yaml.SafeLoader:
+        """
+        Patch the loader to allow = as a value.
+
+        Handles being call multiple times. It only patches the loader on the first call.
+
+        See https://github.com/yaml/pyyaml/issues/89 for more info.
+
+        :return: The patched YAML loader.
+        """
+        if "=" in loader.yaml_implicit_resolvers:
+            loader.yaml_implicit_resolvers.pop("=")
+        return loader
+
     def parse(self, content: str) -> dict:
         """
         Parse the contents of the spec file.
@@ -108,8 +122,12 @@ def parse(self, content: str) -> dict:
         :raises InvalidSpecError: If the spec file is not valid YAML.
         :return: The contents of the spec file as a dict.
         """
+        loader = yaml.SafeLoader
+        patched_loader = self.patch_loader(loader)
         try:
-            return yaml.safe_load(content)
+            return yaml.load(
+                content, Loader=patched_loader  # noqa: S506 Uses a patched safe loader
+            )
         except yaml.YAMLError as e:
             raise picofun.errors.InvalidSpecError() from e
 

tests/test_spec.py

@@ -78,6 +78,12 @@ def test_yaml_spec_parser_invalid_yaml() -> None:
         parser.parse("foo: b: a: r")
 
 
+def test_yaml_spec_parser_value_equals() -> None:
+    """Test the YAMLSpecParser class with a plus symbol as the value."""
+    parser = picofun.spec.YAMLSpecParser()
+    assert parser.parse("key: +") == {"key": "+"}
+
+
 def test_spec_json() -> None:
     """Test the Spec class with JSON."""
     spec = picofun.spec.Spec("tests/data/petstore.json")