Keep unwrap() key order after a value is replaced (#591)

* Keep unwrap() key order after a value is replaced

_replace_at() removes the key from _map and re-inserts it, which moves it
to the end of the dict's insertion order while _body keeps the original
slot. dumps() and keys() walk _body and stay correct; unwrap() walks _map
and does not:

    doc = parse("a = 1\nb = 2\nc = 3\n")
    doc["b"] = 9
    dumps(doc)          # 'a = 1\nb = 9\nc = 3\n'
    list(doc.keys())    # ['a', 'b', 'c']
    list(doc.unwrap())  # ['a', 'c', 'b']

unwrap() moved to _map in #521 on the stated assumption that it "iterates
in the same insertion order as the old self.items()", which a replacement
breaks. Order comes from the body index instead, so the _map fast path is
kept and dumps() stays the reference.

Replacements that genuinely move an element still move: a bare key
promoted to a table has to be emitted after the inline entries, and
unwrap() follows the body there too. Both cases are covered.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: youdie006

tests/test_toml_document.py

@@ -160,6 +160,27 @@ def test_toml_document_without_super_tables() -> None:
     assert "tool" in d
 
 
+def test_unwrap_keeps_key_order_after_replacing_a_value() -> None:
+    # unwrap() reads _map, which re-inserts a replaced key and so moves it
+    # last; the order has to follow the body, like dumps() and keys() do.
+    doc = parse("a = 1\nb = 2\nc = 3\n")
+    doc["b"] = 9
+
+    assert list(doc.unwrap()) == ["a", "b", "c"]
+    assert list(doc.keys()) == ["a", "b", "c"]
+    assert tomlkit.dumps(doc) == "a = 1\nb = 9\nc = 3\n"
+
+
+def test_unwrap_follows_the_body_when_a_value_becomes_a_table() -> None:
+    # here moving the key is correct: a bare key promoted to [table] has to be
+    # emitted after the inline entries, and unwrap() should agree with dumps()
+    doc = parse("a = 1\nb = 2\n")
+    doc["a"] = {"x": 1}
+
+    assert tomlkit.dumps(doc) == "b = 2\n\n[a]\nx = 1\n"
+    assert list(doc.unwrap()) == ["b", "a"]
+
+
 def test_toml_document_unwrap() -> None:
     content = """[tool.poetry]
 name = "foo"

tomlkit/container.py

@@ -67,9 +67,13 @@ def unwrap(self) -> dict[str, Any]:
         # rebuilds a SingleKey from the bare string on every key only to throw
         # it away. Out-of-order keys (a tuple index) still go through
         # OutOfOrderTableProxy so their validation (and fragment merge) runs
-        # exactly as before. _map iterates in the same insertion order as the
-        # old self.items().
-        for key, idx in self._map.items():
+        # exactly as before. _map is keyed for lookup, not ordered: replacing a
+        # value re-inserts its key and moves it last, so take the order from the
+        # body index, which is what dumps() and items() follow.
+        for key, idx in sorted(
+            self._map.items(),
+            key=lambda item: item[1][0] if isinstance(item[1], tuple) else item[1],
+        ):
             if isinstance(idx, tuple):
                 value: Any = OutOfOrderTableProxy(self, idx)
             else: