Add a sort_sets optional keyword argument to dump (swaroopch#42)

bfontaine · Aug 16, 2018 · ea068bd · ea068bd
1 parent ec2c10e
commit ea068bd
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 3 deletions.
diff --git a/edn_format/edn_dump.py b/edn_format/edn_dump.py
@@ -61,12 +61,14 @@ def seq(obj, **kwargs):
 def udump(obj,
           string_encoding=DEFAULT_INPUT_ENCODING,
           keyword_keys=False,
-          sort_keys=False):
+          sort_keys=False,
+          sort_sets=False):
 
     kwargs = {
         "string_encoding": string_encoding,
         "keyword_keys": keyword_keys,
         "sort_keys": sort_keys,
+        "sort_sets": sort_sets,
     }
 
     if obj is None:
@@ -91,6 +93,8 @@ def udump(obj,
     elif isinstance(obj, list):
         return '[{}]'.format(seq(obj, **kwargs))
     elif isinstance(obj, set) or isinstance(obj, frozenset):
+        if sort_sets:
+            obj = sorted(obj)
         return '#{{{}}}'.format(seq(obj, **kwargs))
     elif isinstance(obj, dict) or isinstance(obj, ImmutableDict):
         pairs = obj.items()
@@ -117,11 +121,13 @@ def dump(obj,
          string_encoding=DEFAULT_INPUT_ENCODING,
          output_encoding=DEFAULT_OUTPUT_ENCODING,
          keyword_keys=False,
-         sort_keys=False):
+         sort_keys=False,
+         sort_sets=False):
     outcome = udump(obj,
                     string_encoding=string_encoding,
                     keyword_keys=keyword_keys,
-                    sort_keys=sort_keys)
+                    sort_keys=sort_keys,
+                    sort_sets=sort_sets)
     if __PY3:
         return outcome
     return outcome.encode(output_encoding)
diff --git a/tests.py b/tests.py
@@ -4,6 +4,7 @@
 
 from collections import OrderedDict
 from uuid import uuid4
+import random
 import datetime
 import unittest
 
@@ -130,6 +131,7 @@ def check_roundtrip(self, data_input, **kw):
 
     def test_dump(self):
         self.check_roundtrip({1, 2, 3})
+        self.check_roundtrip({1, 2, 3}, sort_sets=True)
         self.check_roundtrip(
             {Keyword("a"): 1,
              "foo": Keyword("gone"),
@@ -310,6 +312,25 @@ def test_sort_keys(self):
                              {"a": 1, "d": 1, "b": 1, "c": 1},
                              sort_keys=True, keyword_keys=True)
 
+    def test_sort_sets(self):
+        def misordered_set_sequence():
+            """"
+            Return a tuple that, if put in a set then iterated over, doesn't
+            yield elements in the original order
+            """
+            while True:
+                seq = tuple(random.sample(range(10000), random.randint(2, 8)))
+                s = set(seq)
+                if tuple(s) != seq:
+                    return seq
+
+        for _ in range(10):
+            seq = misordered_set_sequence()
+            self.check_roundtrip(set(seq))
+            self.check_dumps("#{{{}}}".format(" ".join(str(i) for i in sorted(seq))),
+                             set(seq),
+                             sort_sets=True)
+
 
 class EdnInstanceTest(unittest.TestCase):
     def test_hashing(self):