From 5721f88aadee66e580b67a3f9086234fb1ef8451 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Eustace?= <sebastien@eustace.io>
Date: Fri, 31 Jul 2020 11:00:53 +0200
Subject: [PATCH] Fix the default ordering of dicts and add support for sorting
 keys

---
 tests/test_items.py | 57 ++++++++++++++++++++++++++++++++++++++++++---
 tomlkit/_compat.py  |  6 +++++
 tomlkit/api.py      |  4 ++--
 tomlkit/items.py    | 16 ++++++++-----
 4 files changed, 72 insertions(+), 11 deletions(-)

diff --git a/tests/test_items.py b/tests/test_items.py
index c4b074ce..5fc30de8 100644
--- a/tests/test_items.py
+++ b/tests/test_items.py
@@ -14,6 +14,7 @@
 from tomlkit import inline_table
 from tomlkit import parse
 from tomlkit._compat import PY2
+from tomlkit._compat import OrderedDict
 from tomlkit.exceptions import NonExistentKey
 from tomlkit.items import Bool
 from tomlkit.items import InlineTable
@@ -208,8 +209,56 @@ def test_dicts_are_converted_to_tables():
     )
 
 
-def test_dicts_are_converted_to_tables_and_sorted():
-    t = item({"foo": {"bar": "baz", "abc": 123, "baz": [{"c": 3, "b": 2, "a": 1}]}})
+def test_dicts_are_converted_to_tables_and_keep_order():
+    t = item(
+        OrderedDict(
+            [
+                (
+                    "foo",
+                    OrderedDict(
+                        [
+                            ("bar", "baz"),
+                            ("abc", 123),
+                            ("baz", [OrderedDict([("c", 3), ("b", 2), ("a", 1)])]),
+                        ]
+                    ),
+                )
+            ]
+        )
+    )
+
+    assert (
+        t.as_string()
+        == """[foo]
+bar = "baz"
+abc = 123
+
+[[foo.baz]]
+c = 3
+b = 2
+a = 1
+"""
+    )
+
+
+def test_dicts_are_converted_to_tables_and_are_sorted_if_requested():
+    t = item(
+        OrderedDict(
+            [
+                (
+                    "foo",
+                    OrderedDict(
+                        [
+                            ("bar", "baz"),
+                            ("abc", 123),
+                            ("baz", [OrderedDict([("c", 3), ("b", 2), ("a", 1)])]),
+                        ]
+                    ),
+                )
+            ]
+        ),
+        _sort_keys=True,
+    )
 
     assert (
         t.as_string()
@@ -226,7 +275,9 @@ def test_dicts_are_converted_to_tables_and_sorted():
 
 
 def test_dicts_with_sub_dicts_are_properly_converted():
-    t = item({"foo": {"bar": {"string": "baz"}, "int": 34, "float": 3.14}})
+    t = item(
+        {"foo": {"bar": {"string": "baz"}, "int": 34, "float": 3.14}}, _sort_keys=True
+    )
 
     assert (
         t.as_string()
diff --git a/tomlkit/_compat.py b/tomlkit/_compat.py
index d5cb1f8e..8d3b0ae3 100644
--- a/tomlkit/_compat.py
+++ b/tomlkit/_compat.py
@@ -150,6 +150,12 @@ def _name_from_offset(delta):
     long = int
 
 
+if PY36:
+    OrderedDict = dict
+else:
+    from collections import OrderedDict
+
+
 def decode(string, encodings=None):
     if not PY2 and not isinstance(string, bytes):
         return string
diff --git a/tomlkit/api.py b/tomlkit/api.py
index 3f475263..3de41219 100644
--- a/tomlkit/api.py
+++ b/tomlkit/api.py
@@ -34,12 +34,12 @@ def loads(string):  # type: (str) -> _TOMLDocument
     return parse(string)
 
 
-def dumps(data):  # type: (_TOMLDocument) -> str
+def dumps(data, sort_keys=False):  # type: (_TOMLDocument, bool) -> str
     """
     Dumps a TOMLDocument into a string.
     """
     if not isinstance(data, _TOMLDocument) and isinstance(data, dict):
-        data = item(data)
+        data = item(data, _sort_keys=sort_keys)
 
     return data.as_string()
 
diff --git a/tomlkit/items.py b/tomlkit/items.py
index 8218f9cd..184ffe7d 100644
--- a/tomlkit/items.py
+++ b/tomlkit/items.py
@@ -28,7 +28,7 @@
     from functools import lru_cache
 
 
-def item(value, _parent=None):
+def item(value, _parent=None, _sort_keys=False):
     from .container import Container
 
     if isinstance(value, Item):
@@ -42,8 +42,11 @@ def item(value, _parent=None):
         return Float(value, Trivia(), str(value))
     elif isinstance(value, dict):
         val = Table(Container(), Trivia(), False)
-        for k, v in sorted(value.items(), key=lambda i: (isinstance(i[1], dict), i[0])):
-            val[k] = item(v, _parent=val)
+        for k, v in sorted(
+            value.items(),
+            key=lambda i: (isinstance(i[1], dict), i[0] if _sort_keys else 1),
+        ):
+            val[k] = item(v, _parent=val, _sort_keys=_sort_keys)
 
         return val
     elif isinstance(value, list):
@@ -57,13 +60,14 @@ def item(value, _parent=None):
                 table = Table(Container(), Trivia(), True)
 
                 for k, _v in sorted(
-                    v.items(), key=lambda i: (isinstance(i[1], dict), i[0])
+                    v.items(),
+                    key=lambda i: (isinstance(i[1], dict), i[0] if _sort_keys else 1),
                 ):
-                    i = item(_v)
+                    i = item(_v, _sort_keys=_sort_keys)
                     if isinstance(table, InlineTable):
                         i.trivia.trail = ""
 
-                    table[k] = item(i)
+                    table[k] = item(i, _sort_keys=_sort_keys)
 
                 v = table