From 6394366c349d9aed9e86143efaf3a5aa69f10ec4 Mon Sep 17 00:00:00 2001
From: Victor Stinner <haypo@maggie.(none)>
Date: Tue, 22 Jun 2010 20:33:19 +0200
Subject: [PATCH] RestorableDict(): support creation of a new key

---
 sandbox/restorable_dict.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/sandbox/restorable_dict.py b/sandbox/restorable_dict.py
index e223af7..f655c6e 100644
--- a/sandbox/restorable_dict.py
+++ b/sandbox/restorable_dict.py
@@ -2,10 +2,14 @@ class RestorableDict:
     def __init__(self, dict):
         self.dict = dict
         self.original = {}
+        self.delete = set()
 
     def __setitem__(self, key, value):
-        if key not in self.original:
-            self.original[key] = self.dict[key]
+        if (key not in self.original) and (key not in self.delete):
+            if key in self.dict:
+                self.original[key] = self.dict[key]
+            else:
+                self.delete.add(key)
         self.dict[key] = value
 
     def __delitem__(self, key):
@@ -15,6 +19,8 @@ def copy(self):
         return self.dict.copy()
 
     def restore(self):
+        for key in self.delete:
+            del self.dict[key]
         self.dict.update(self.original)
         self.original.clear()