Change Map.delete(key): raise a KeyError if key is not in the map

Author: 1st1

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

immutables/map.py

@@ -126,7 +126,7 @@ def find(self, shift, hash, key):
         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 @@ def find(self, shift, hash, key):
         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 @@ def delete(self, key):
         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

tests/test_map.py

@@ -269,9 +269,10 @@ def test_map_stress(self):
             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 @@ def test_map_delete_1(self):
         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 @@ def test_map_delete_2(self):
         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 @@ def test_map_delete_2(self):
         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 @@ def test_map_delete_5(self):
 
         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):