From b057a59bb909c0222ab86ff9922e80e7466731ae Mon Sep 17 00:00:00 2001
From: wim glenn <hey@wimglenn.com>
Date: Wed, 27 Jun 2018 00:11:39 -0500
Subject: [PATCH] support PyYAML 4.1.0

---
 .travis.yml   |  5 +++++
 oyaml.py      | 17 ++++++++++++-----
 setup.py      |  2 +-
 test_oyaml.py | 16 +++++++++++++---
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 1a9122a..ffcf965 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -11,12 +11,17 @@ python:
   - "pypy"
   - "pypy3.5"
 
+env:
+  - PYYAML_VERSION="3.12"
+  - PYYAML_VERSION="4.1"
+
 matrix:
   fast_finish: true
   allow_failures:
   - python: "nightly"
 
 install: 
+  - pip install pyyaml~=$PYYAML_VERSION
   - pip install --editable .
   - pip install pytest-cov coveralls
 
diff --git a/oyaml.py b/oyaml.py
index 0fdf8e4..b78ebde 100644
--- a/oyaml.py
+++ b/oyaml.py
@@ -15,11 +15,18 @@ def map_constructor(loader, node):
     loader.flatten_mapping(node)
     return OrderedDict(loader.construct_pairs(node))
 
-
-pyyaml.add_representer(dict, map_representer)
-pyyaml.add_representer(OrderedDict, map_representer)
-pyyaml.add_representer(dict, map_representer, Dumper=pyyaml.dumper.SafeDumper)
-pyyaml.add_representer(OrderedDict, map_representer, Dumper=pyyaml.dumper.SafeDumper)
+if pyyaml.safe_dump is pyyaml.dump:
+    # PyYAML >= 4
+    SafeDumper = pyyaml.dumper.Dumper
+    DangerDumper = pyyaml.dumper.DangerDumper
+else:
+    SafeDumper = pyyaml.dumper.SafeDumper
+    DangerDumper = pyyaml.dumper.Dumper
+
+pyyaml.add_representer(dict, map_representer, Dumper=SafeDumper)
+pyyaml.add_representer(OrderedDict, map_representer, Dumper=SafeDumper)
+pyyaml.add_representer(dict, map_representer, Dumper=DangerDumper)
+pyyaml.add_representer(OrderedDict, map_representer, Dumper=DangerDumper)
 
 
 if sys.version_info < (3, 7):
diff --git a/setup.py b/setup.py
index 0d8f6fb..df039dc 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name='oyaml',
-    version='0.4',
+    version='0.5',
     description='Ordered YAML: drop-in replacement for PyYAML which preserves dict ordering',
     long_description=open('README.rst').read(),
     author='Wim Glenn',
diff --git a/test_oyaml.py b/test_oyaml.py
index 69ff6d1..260d317 100644
--- a/test_oyaml.py
+++ b/test_oyaml.py
@@ -58,11 +58,12 @@ def test_loads_to_std_dict():
     assert isinstance(loaded, dict)
 
 
-def test_subclass_dump():
+class MyOrderedDict(OrderedDict):
+    pass
 
-    class MyOrderedDict(OrderedDict):
-        pass
 
+@pytest.mark.skipif(yaml.pyyaml.__version__ >= '4', reason="requires PyYAML version < 4")
+def test_subclass_dump_pyyaml3():
     data = MyOrderedDict([('x', 1), ('y', 2)])
     assert '!!python/object/apply:test_oyaml.MyOrderedDict' in yaml.dump(data)
     with pytest.raises(yaml.pyyaml.representer.RepresenterError) as cm:
@@ -70,6 +71,15 @@ class MyOrderedDict(OrderedDict):
     assert str(cm.value) == "cannot represent an object: MyOrderedDict([('x', 1), ('y', 2)])"
 
 
+@pytest.mark.skipif(yaml.pyyaml.__version__ < '4', reason="requires PyYAML version >= 4")
+def test_subclass_dump_pyyaml4():
+    data = MyOrderedDict([('x', 1), ('y', 2)])
+    assert '!!python/object/apply:test_oyaml.MyOrderedDict' in yaml.danger_dump(data)
+    with pytest.raises(yaml.pyyaml.representer.RepresenterError) as cm:
+        yaml.dump(data)
+    assert str(cm.value) == "('cannot represent an object', MyOrderedDict([('x', 1), ('y', 2)]))"
+
+
 def test_anchors_and_references():
     text = '''
         defaults: