Repository: awslabs/aws-cfn-template-flip

Issue:
yaml ordereddict fix breaks when yaml is a .egg
The fix for #14 doesn't work in all cases. When PyYAML is installed in a .egg file the load of yaml as custom_yaml fails. For troposphere, [here](https://travis-ci.org/cloudtools/troposphere/jobs/256102858) is an example of the tests failing. The issue is the imp module does not know how how to handle hooks to zipimport. I have yet to find a good alternate solution to this [code](https://github.com/awslabs/aws-cfn-template-flip/blob/master/cfn_flip/custom_yaml.py#L15).

Discussion:

Patch:
diff --git a/cfn_flip/__init__.py b/cfn_flip/__init__.py
index 34e9257..d5b112a 100644
--- a/cfn_flip/__init__.py
+++ b/cfn_flip/__init__.py
@@ -10,19 +10,11 @@ or in the "license" file accompanying this file. This file is distributed on an
 
 from .clean import clean
 from .custom_json import DateTimeAwareJsonEncoder
-from .custom_yaml import custom_yaml
+from .custom_yaml import CustomDumper, CustomLoader
 import collections
 import json
+import yaml
 
-class MyDumper(custom_yaml.Dumper):
-  """
-  Indent block sequences from parent using more common style
-  ("  - entry"  vs "- entry").  
-  Causes fewer problems with validation and tools.
-  """
-  
-  def increase_indent(self,flow=False, indentless=False):
-    return super(MyDumper,self).increase_indent(flow, False)
 
 def to_json(template, clean_up=False):
     """
@@ -30,7 +22,7 @@ def to_json(template, clean_up=False):
     undoing yaml short syntax where detected
     """
 
-    data = custom_yaml.load(template)
+    data = yaml.load(template, Loader=CustomLoader)
 
     if clean_up:
         data = clean(data)
@@ -48,7 +40,7 @@ def to_yaml(template, clean_up=False):
     if clean_up:
         data = clean(data)
 
-    return custom_yaml.dump(data, Dumper=MyDumper, default_flow_style=False)
+    return yaml.dump(data, Dumper=CustomDumper, default_flow_style=False)
 
 def flip(template, clean_up=False):
     """
diff --git a/cfn_flip/custom_yaml.py b/cfn_flip/custom_yaml.py
index a2ff89c..b1ea022 100644
--- a/cfn_flip/custom_yaml.py
+++ b/cfn_flip/custom_yaml.py
@@ -8,16 +8,30 @@ Licensed under the Apache License, Version 2.0 (the "License"). You may not use
 or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.                                                 
 """
 
-import imp
 import six
 import collections
+import yaml
 
-custom_yaml = imp.load_module("custom_yaml", *imp.find_module("yaml"))
 
 TAG_MAP = "tag:yaml.org,2002:map"
 TAG_STRING = "tag:yaml.org,2002:str"
 UNCONVERTED_SUFFIXES = ["Ref", "Condition"]
 
+class CustomDumper(yaml.Dumper):
+  """
+  Indent block sequences from parent using more common style
+  ("  - entry"  vs "- entry").
+  Causes fewer problems with validation and tools.
+  """
+
+  def increase_indent(self,flow=False, indentless=False):
+    return super(CustomDumper,self).increase_indent(flow, False)
+
+
+class CustomLoader(yaml.Loader):
+    pass
+
+
 def multi_constructor(loader, tag_suffix, node):
     """
     Deal with !Ref style function format
@@ -30,11 +44,11 @@ def multi_constructor(loader, tag_suffix, node):
 
     if tag_suffix == "Fn::GetAtt":
         constructor = construct_getatt
-    elif isinstance(node, custom_yaml.ScalarNode):
+    elif isinstance(node, yaml.ScalarNode):
         constructor = loader.construct_scalar
-    elif isinstance(node, custom_yaml.SequenceNode):
+    elif isinstance(node, yaml.SequenceNode):
         constructor = loader.construct_sequence
-    elif isinstance(node, custom_yaml.MappingNode):
+    elif isinstance(node, yaml.MappingNode):
         constructor = loader.construct_mapping
     else:
         raise "Bad tag: !{}".format(tag_suffix)
@@ -126,8 +140,8 @@ def representer(dumper, data):
     return dumper.represent_scalar(tag, data)
 
 # Customise our yaml
-custom_yaml.add_representer(six.text_type, lambda dumper, value: dumper.represent_scalar(TAG_STRING, value))
-custom_yaml.add_constructor(TAG_MAP, construct_mapping)
-custom_yaml.add_multi_constructor("!", multi_constructor)
-custom_yaml.add_representer(collections.OrderedDict, representer)
-custom_yaml.add_representer(dict, representer)
+CustomDumper.add_representer(six.text_type, lambda dumper, value: dumper.represent_scalar(TAG_STRING, value))
+CustomLoader.add_constructor(TAG_MAP, construct_mapping)
+CustomLoader.add_multi_constructor("!", multi_constructor)
+CustomDumper.add_representer(collections.OrderedDict, representer)
+CustomDumper.add_representer(dict, representer)
diff --git a/setup.py b/setup.py
index 52de37c..e6acc7c 100644
--- a/setup.py
+++ b/setup.py
@@ -24,6 +24,7 @@ setup(
         "six",
     ],
     zip_safe=False,
+    test_suite="tests",
     entry_points={
         "console_scripts": ["cfn-flip=cfn_flip.main:main"],
     },


Test patch:
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_flip.py b/tests/test_flip.py
index 17d7236..fc50b67 100644
--- a/tests/test_flip.py
+++ b/tests/test_flip.py
@@ -9,9 +9,10 @@ or in the "license" file accompanying this file. This file is distributed on an
 """
 
 import cfn_flip
+from cfn_flip.custom_yaml import CustomLoader
 import json
 import unittest
-from cfn_flip.custom_yaml import custom_yaml
+import yaml
 
 
 class CfnFlipTestCase(unittest.TestCase):
@@ -33,10 +34,10 @@ class CfnFlipTestCase(unittest.TestCase):
             self.clean_yaml = f.read()
 
         self.parsed_json = json.loads(self.input_json)
-        self.parsed_yaml = custom_yaml.load(self.input_yaml)
+        self.parsed_yaml = yaml.load(self.input_yaml, Loader=CustomLoader)
 
         self.parsed_clean_json = json.loads(self.clean_json)
-        self.parsed_clean_yaml = custom_yaml.load(self.clean_yaml)
+        self.parsed_clean_yaml = yaml.load(self.clean_yaml, Loader=CustomLoader)
 
         self.bad_data = "<!DOCTYPE html>\n\n<html>\n\tThis isn't right!\n</html>"
 
@@ -76,7 +77,7 @@ class CfnFlipTestCase(unittest.TestCase):
         with self.assertRaises(ValueError):
             json.loads(actual)
 
-        parsed_actual = custom_yaml.load(actual)
+        parsed_actual = yaml.load(actual, Loader=CustomLoader)
 
         self.assertDictEqual(parsed_actual, self.parsed_yaml)
 
@@ -111,7 +112,7 @@ class CfnFlipTestCase(unittest.TestCase):
         with self.assertRaises(ValueError):
             json.loads(actual)
 
-        parsed_actual = custom_yaml.load(actual)
+        parsed_actual = yaml.load(actual, Loader=CustomLoader)
 
         self.assertDictEqual(parsed_actual, self.parsed_yaml)
 
@@ -139,7 +140,7 @@ class CfnFlipTestCase(unittest.TestCase):
         with self.assertRaises(ValueError):
             json.loads(actual)
 
-        parsed_actual = custom_yaml.load(actual)
+        parsed_actual = yaml.load(actual, Loader=CustomLoader)
 
         self.assertDictEqual(parsed_actual, self.parsed_clean_yaml)
 
diff --git a/tests/test_yaml_patching.py b/tests/test_yaml_patching.py
index 5079357..db23442 100644
--- a/tests/test_yaml_patching.py
+++ b/tests/test_yaml_patching.py
@@ -9,6 +9,7 @@ or in the "license" file accompanying this file. This file is distributed on an
 """
 
 import cfn_flip
+import collections
 import json
 import unittest
 import yaml
@@ -19,7 +20,7 @@ class YamlPatchTestCase(unittest.TestCase):
     Check that we don't patch yaml for everybody
     """
 
-    def test_yaml_ordered_dict(self):
+    def test_yaml_no_ordered_dict(self):
         """
         cfn-flip patches yaml to use OrderedDict by default
         Check that we don't do this for folks who import cfn_flip and yaml
@@ -29,3 +30,14 @@ class YamlPatchTestCase(unittest.TestCase):
         data = yaml.load(yaml_string)
 
         self.assertEqual(type(data), dict)
+
+    def test_yaml_no_ordered_dict(self):
+        """
+        cfn-flip patches yaml to use OrderedDict by default
+        Check that we do this for normal cfn_flip use cases
+        """
+
+        yaml_string = "key: value"
+        data = yaml.load(yaml_string, Loader=cfn_flip.CustomLoader)
+
+        self.assertEqual(type(data), collections.OrderedDict)