Merge branch 'yaml-custom-tags'

PR #1434

* yaml-custom-tags:
  Add support for parsing custom yaml tags used in CFN

chalice/package.py

@@ -4,10 +4,12 @@
 import json
 import os
 
+import six
 from typing import Any, Optional, Dict, List, Set, Union  # noqa
 from typing import cast
 import yaml
 from yaml.scanner import ScannerError
+from yaml.nodes import Node, ScalarNode, SequenceNode
 
 from chalice.deploy.swagger import (
     CFNSwaggerGenerator, TerraformSwaggerGenerator)
@@ -1121,10 +1123,12 @@ def is_yaml_template(cls, template_name):
 
     def serialize_template(self, contents):
         # type: (Dict[str, Any]) -> str
-        return yaml.dump(contents, allow_unicode=True)
+        return yaml.safe_dump(contents, allow_unicode=True)
 
     def load_template(self, file_contents, filename=''):
         # type: (str, str) -> Dict[str, Any]
+        yaml.SafeLoader.add_multi_constructor(
+            tag_prefix='!', multi_constructor=self._custom_sam_instrinsics)
         try:
             return yaml.load(
                 file_contents,
@@ -1133,3 +1137,24 @@ def load_template(self, file_contents, filename=''):
         except ScannerError:
             raise RuntimeError(
                 'Expected %s to be valid YAML template.' % filename)
+
+    def _custom_sam_instrinsics(self, loader, tag_prefix, node):
+        # type: (yaml.SafeLoader, str, Node) -> Dict[str, Any]
+        tag = node.tag[1:]
+        if tag not in ['Ref', 'Condition']:
+            tag = 'Fn::%s' % tag
+        value = self._get_value(loader, node)
+        return {tag: value}
+
+    def _get_value(self, loader, node):
+        # type: (yaml.SafeLoader, Node) -> Any
+        if node.tag[1:] == 'GetAtt' and isinstance(node.value,
+                                                   six.string_types):
+            value = node.value.split('.', 1)
+        elif isinstance(node, ScalarNode):
+            value = loader.construct_scalar(node)
+        elif isinstance(node, SequenceNode):
+            value = loader.construct_sequence(node)
+        else:
+            value = loader.construct_mapping(node)
+        return value

tests/unit/test_package.py

@@ -1397,3 +1397,26 @@ def test_merge_can_change_type(self):
 ])
 def test_to_cfn_resource_name(filename, is_yaml):
     assert package.YAMLTemplateSerializer.is_yaml_template(filename) == is_yaml
+
+
+@pytest.mark.parametrize('yaml_contents,expected', [
+    ('foo: bar', {'foo': 'bar'}),
+    ('foo: !Ref bar', {'foo': {'Ref': 'bar'}}),
+    ('foo: !GetAtt Bar.Baz', {'foo': {'Fn::GetAtt': ['Bar', 'Baz']}}),
+    ('foo: !FooBar [!Baz YetAnother, "hello"]',
+     {'foo': {'Fn::FooBar': [{'Fn::Baz': 'YetAnother'}, 'hello']}}),
+    ('foo: !SomeTag {"a": "1"}', {'foo': {'Fn::SomeTag': {'a': '1'}}}),
+    ('foo: !GetAtt Foo.Bar.Baz', {'foo': {'Fn::GetAtt': ['Foo', 'Bar.Baz']}}),
+    ('foo: !Condition Other', {'foo': {'Condition': 'Other'}}),
+    ('foo: !GetAtt ["a", "b"]', {'foo': {'Fn::GetAtt': ['a', 'b']}}),
+])
+def test_supports_custom_tags(yaml_contents, expected):
+    serialize = package.YAMLTemplateSerializer()
+    actual = serialize.load_template(yaml_contents)
+    assert actual == expected
+    # Also verify we can serialize then parse them back to what we originally
+    # loaded.  Note that this is not the same thing as round tripping as
+    # we convert things like '!Ref foo' over to {'Ref': 'foo'}.
+    yaml_str = serialize.serialize_template(actual).strip()
+    reparsed = serialize.load_template(yaml_str)
+    assert actual == reparsed