From 264dbf77ab5797536605653dae48d2453b450670 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Eustace?= <sebastien@eustace.io>
Date: Fri, 24 Jan 2020 16:17:40 +0100
Subject: [PATCH] Fix behavior for out of order tables

---
 tests/test_toml_document.py | 62 +++++++++++++++++++++++++
 tomlkit/container.py        | 92 +++++++++++++++++++++++++++++++------
 2 files changed, 141 insertions(+), 13 deletions(-)

diff --git a/tests/test_toml_document.py b/tests/test_toml_document.py
index 28b55808..87f4c741 100644
--- a/tests/test_toml_document.py
+++ b/tests/test_toml_document.py
@@ -7,9 +7,11 @@
 
 from datetime import datetime
 
+import pytest
 import tomlkit
 from tomlkit import parse
 from tomlkit._utils import _utc
+from tomlkit.exceptions import NonExistentKey
 
 
 def test_document_is_a_dict(example):
@@ -473,3 +475,63 @@ def test_declare_sub_table_with_intermediate_table():
     doc = parse(content)
     assert {"tommy": 87, "mary": 66, "bob": {"score": 91}} == doc["students"]
     assert {"tommy": 87, "mary": 66, "bob": {"score": 91}} == doc.get("students")
+
+
+def test_values_can_still_be_set_for_out_of_order_tables():
+    content = """
+[a.a]
+key = "value"
+
+[a.b]
+
+[a.a.c]
+"""
+
+    doc = parse(content)
+    doc["a"]["a"]["key"] = "new_value"
+
+    assert "new_value" == doc["a"]["a"]["key"]
+
+    expected = """
+[a.a]
+key = "new_value"
+
+[a.b]
+
+[a.a.c]
+"""
+
+    assert expected == doc.as_string()
+
+    doc["a"]["a"]["bar"] = "baz"
+
+    expected = """
+[a.a]
+key = "new_value"
+bar = "baz"
+
+[a.b]
+
+[a.a.c]
+"""
+
+    assert expected == doc.as_string()
+
+    del doc["a"]["a"]["key"]
+
+    expected = """
+[a.a]
+bar = "baz"
+
+[a.b]
+
+[a.a.c]
+"""
+
+    assert expected == doc.as_string()
+
+    with pytest.raises(NonExistentKey):
+        doc["a"]["a"]["key"]
+
+    with pytest.raises(NonExistentKey):
+        del doc["a"]["a"]["key"]
diff --git a/tomlkit/container.py b/tomlkit/container.py
index 80d9f139..9838dc1a 100644
--- a/tomlkit/container.py
+++ b/tomlkit/container.py
@@ -524,21 +524,12 @@ def __getitem__(self, key):  # type: (Union[Key, str]) -> Union[Item, Container]
             raise NonExistentKey(key)
 
         if isinstance(idx, tuple):
-            container = Container(True)
-
-            for i in idx:
-                item = self._body[i][1]
-
-                if isinstance(item, Table):
-                    for k, v in item.value.body:
-                        container.append(k, v)
-                else:
-                    container.append(key, item)
-
-            return container
+            # The item we are getting is an out of order table
+            # so we need a proxy to retrieve the proper objects
+            # from the parent container
+            return OutOfOrderTableProxy(self, idx)
 
         item = self._body[idx][1]
-
         if item.is_boolean():
             return item.value
 
@@ -571,6 +562,9 @@ def _replace(
     def _replace_at(
         self, idx, new_key, value
     ):  # type: (Union[int, Tuple[int]], Union[Key, str], Item) -> None
+        if not isinstance(new_key, Key):
+            new_key = Key(new_key)
+
         if isinstance(idx, tuple):
             for i in idx[1:]:
                 self._body[i] = (None, Null())
@@ -580,6 +574,8 @@ def _replace_at(
         k, v = self._body[idx]
 
         self._map[new_key] = self._map.pop(k)
+        if new_key != k:
+            super(Container, self).__delitem__(k)
 
         if isinstance(self._map[new_key], tuple):
             self._map[new_key] = self._map[new_key][0]
@@ -640,3 +636,73 @@ def __copy__(self):  # type: () -> Container
         c._map.update(self._map)
 
         return c
+
+
+class OutOfOrderTableProxy(object):
+    def __init__(self, container, indices):  # type: (Container, Tuple) -> None
+        self._container = container
+        self._internal_container = Container(self._container.parsing)
+        self._tables = []
+        self._tables_map = {}
+        self._map = {}
+
+        for i in indices:
+            key, item = self._container._body[i]
+
+            if isinstance(item, Table):
+                self._tables.append(item)
+                table_idx = len(self._tables) - 1
+                for k, v in item.value.body:
+                    self._internal_container.append(k, v)
+                    self._tables_map[k] = table_idx
+            else:
+                self._internal_container.append(key, item)
+                self._map[key] = i
+
+    def __getitem__(self, key):  # type: (Union[Key, str]) -> Any
+        if key not in self._internal_container:
+            raise NonExistentKey(key)
+
+        return self._internal_container[key]
+
+    def __setitem__(self, key, item):  # type: (Union[Key, str], Any) -> None
+        if key in self._map:
+            idx = self._map[key]
+            self._container._replace_at(idx, key, item)
+        elif key in self._tables_map:
+            table = self._tables[self._tables_map[key]]
+            table[key] = item
+        elif self._tables:
+            table = self._tables[0]
+            table[key] = item
+        else:
+            self._container[key] = item
+
+    def __delitem__(self, key):  # type: (Union[Key, str]) -> None
+        if key in self._map:
+            idx = self._map[key]
+            del self._container[key]
+            del self._map[key]
+        elif key in self._tables_map:
+            table = self._tables[self._tables_map[key]]
+            del table[key]
+            del self._tables_map[key]
+        else:
+            raise NonExistentKey(key)
+
+        del self._internal_container[key]
+
+    def __str__(self):
+        return str(self._internal_container)
+
+    def __repr__(self):
+        return repr(self._internal_container)
+
+    def __eq__(self, other):  # type: (Dict) -> bool
+        if not isinstance(other, dict):
+            return NotImplemented
+
+        return self._internal_container == other
+
+    def __getattr__(self, attribute):
+        return getattr(self._internal_container, attribute)