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()