From 5804f0c47e60c7c2a0915cce2dbcfc7593881761 Mon Sep 17 00:00:00 2001
From: Yury Selivanov <yury@magic.io>
Date: Tue, 3 Apr 2018 00:01:41 -0400
Subject: Change Map.delete(key): raise a KeyError if key is not in the map

---
 immutables/_map.c |  4 ++--
 immutables/map.py |  7 +++----
 tests/test_map.py | 27 ++++++++++++++++++---------
 3 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/immutables/_map.c b/immutables/_map.c
index 445b713..3ed0393 100644
--- a/immutables/_map.c
+++ b/immutables/_map.c
@@ -2355,8 +2355,8 @@ map_without(MapObject *o, PyObject *key)
         case W_EMPTY:
             return map_new();
         case W_NOT_FOUND:
-            Py_INCREF(o);
-            return o;
+            PyErr_SetObject(PyExc_KeyError, key);
+            return NULL;
         case W_NEWNODE: {
             assert(new_root != NULL);
 
diff --git a/immutables/map.py b/immutables/map.py
index ffed188..b6d55dc 100644
--- a/immutables/map.py
+++ b/immutables/map.py
@@ -126,7 +126,7 @@ class BitmapNode:
         if key == key_or_null:
             return val_or_node
 
-        raise KeyError
+        raise KeyError(key)
 
     def without(self, shift, hash, key):
         bit = map_bitpos(hash, shift)
@@ -240,7 +240,7 @@ class CollisionNode:
         for i in range(0, self.size, 2):
             if self.array[i] == key:
                 return self.array[i + 1]
-        raise KeyError
+        raise KeyError(key)
 
     def assoc(self, shift, hash, key, val, added_leaf):
         if hash == self.hash:
@@ -388,8 +388,7 @@ class Map:
         if res is W_EMPTY:
             return Map()
         elif res is W_NOT_FOUND:
-            # raise KeyError(key)
-            return self
+            raise KeyError(key)
         else:
             m = Map.__new__(Map)
             m.__count = self.__count - 1
diff --git a/tests/test_map.py b/tests/test_map.py
index 03dc9a0..3e7633a 100644
--- a/tests/test_map.py
+++ b/tests/test_map.py
@@ -269,9 +269,10 @@ class BaseMapTest:
             self.assertEqual(len(dm), len(hm))
 
             for i, key in enumerate(keys_to_delete):
-                hm = hm.delete(str(key))
+                if str(key) in dm:
+                    hm = hm.delete(str(key))
+                    dm.pop(str(key))
                 self.assertEqual(hm.get(str(key), 'not found'), 'not found')
-                dm.pop(str(key), None)
                 self.assertEqual(len(d), len(h))
 
                 if not (i % TEST_ITERS_EVERY):
@@ -317,8 +318,9 @@ class BaseMapTest:
         h = h.delete(D)
         self.assertEqual(len(h), orig_len - 2)
 
-        h2 = h.delete(Z)
-        self.assertIs(h2, h)
+        with self.assertRaises(KeyError) as ex:
+            h.delete(Z)
+        self.assertIs(ex.exception.args[0], Z)
 
         h = h.delete(A)
         self.assertEqual(len(h), orig_len - 3)
@@ -358,7 +360,9 @@ class BaseMapTest:
         with self.assertRaisesRegex(ValueError, 'cannot compare'):
             h.delete(Er)
 
-        h = h.delete(Z)
+        with self.assertRaises(KeyError) as ex:
+            h.delete(Z)
+        self.assertIs(ex.exception.args[0], Z)
         self.assertEqual(len(h), orig_len)
 
         h = h.delete(C)
@@ -373,8 +377,10 @@ class BaseMapTest:
         self.assertEqual(h.get(D), 'd')
         self.assertEqual(h.get(E), 'e')
 
-        h = h.delete(A)
-        h = h.delete(B)
+        with self.assertRaises(KeyError):
+            h = h.delete(A)
+        with self.assertRaises(KeyError):
+            h = h.delete(B)
         h = h.delete(D)
         h = h.delete(E)
         self.assertEqual(len(h), 0)
@@ -499,11 +505,14 @@ class BaseMapTest:
 
         h = h.delete(keys[1])
         self.assertEqual(len(h), 14)
-        h = h.delete(keys[1])
+        with self.assertRaises(KeyError) as ex:
+            h.delete(keys[1])
+        self.assertIs(ex.exception.args[0], keys[1])
         self.assertEqual(len(h), 14)
 
         for key in keys:
-            h = h.delete(key)
+            if key in h:
+                h = h.delete(key)
         self.assertEqual(len(h), 0)
 
     def test_map_items_1(self):
-- 
cgit v1.2.3