[pypy-commit] pypy py3k: Implement rdict.pop(). Always with default.

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: py3k
Changeset: r48162:5364600194fb
Date: 2011-10-17 20:10 +0200
http://bitbucket.org/pypy/pypy/changeset/5364600194fb/

Log:	Implement rdict.pop(). Always with default.

diff --git a/pypy/annotation/unaryop.py b/pypy/annotation/unaryop.py
--- a/pypy/annotation/unaryop.py
+++ b/pypy/annotation/unaryop.py
@@ -407,6 +407,7 @@
         return dct.dictdef.read_value()
 
     method_setdefault = method_get
+    method_pop = method_get
 
     def method_copy(dct):
         return SomeDict(dct.dictdef)
diff --git a/pypy/rpython/lltypesystem/rdict.py b/pypy/rpython/lltypesystem/rdict.py
--- a/pypy/rpython/lltypesystem/rdict.py
+++ b/pypy/rpython/lltypesystem/rdict.py
@@ -275,6 +275,13 @@
         v_res = hop.gendirectcall(ll_setdefault, v_dict, v_key, v_default)
         return self.recast_value(hop.llops, v_res)
 
+    def rtype_method_pop(self, hop):
+        v_dict, v_key, v_default = hop.inputargs(self, self.key_repr,
+                                                 self.value_repr)
+        hop.exception_cannot_occur()
+        v_res = hop.gendirectcall(ll_dict_pop, v_dict, v_key, v_default)
+        return self.recast_value(hop.llops, v_res)
+
     def rtype_method_copy(self, hop):
         v_dict, = hop.inputargs(self)
         hop.exception_cannot_occur()
@@ -492,6 +499,15 @@
         raise KeyError
     _ll_dict_del(d, i)
 
+def ll_dict_pop(d, key, default):
+    i = ll_dict_lookup(d, key, d.keyhash(key))
+    if not i & HIGHEST_BIT:
+        value = ll_get_value(d, i)
+        _ll_dict_del(d, i)
+        return value
+    else:
+        return default
+
 @jit.dont_look_inside
 def _ll_dict_del(d, i):
     d.entries.mark_deleted(i)
diff --git a/pypy/rpython/ootypesystem/rdict.py b/pypy/rpython/ootypesystem/rdict.py
--- a/pypy/rpython/ootypesystem/rdict.py
+++ b/pypy/rpython/ootypesystem/rdict.py
@@ -105,6 +105,13 @@
         v_res = hop.gendirectcall(ll_dict_setdefault, v_dict, v_key, v_default)
         return self.recast_value(hop.llops, v_res)
 
+    def rtype_method_pop(self, hop):
+        v_dict, v_key, v_default = hop.inputargs(self, self.key_repr,
+                                                 self.value_repr)
+        hop.exception_cannot_occur()
+        v_res = hop.gendirectcall(ll_dict_pop, v_dict, v_key, v_default)
+        return self.recast_value(hop.llops, v_res)
+
     def rtype_method_copy(self, hop):
         v_dict, = hop.inputargs(self)
         cDICT = hop.inputconst(ootype.Void, self.lowleveltype)
@@ -334,6 +341,14 @@
         d.ll_set(key, default)
         return default
 
+def ll_dict_pop(d, key, default):
+    if d.ll_contains(key):
+        value = d.ll_get(key)
+        d.ll_remove(key)
+        return value
+    else:
+        return default
+
 def _make_ll_keys_values_items(kind):
     def ll_dict_kvi(LIST, d):
         length = d.ll_length()
diff --git a/pypy/rpython/test/test_rdict.py b/pypy/rpython/test/test_rdict.py
--- a/pypy/rpython/test/test_rdict.py
+++ b/pypy/rpython/test/test_rdict.py
@@ -171,6 +171,16 @@
         res = self.interpret(func, ())
         assert res == 422
 
+    def test_dict_pop(self):
+        def func():
+            dic = {}
+            x1 = dic.pop('hi', 42)
+            dic['blah'] = 1
+            x2 = dic.pop('blah', 2)
+            return len(dic) * 100 + x1 * 10 + x2
+        res = self.interpret(func, ())
+        assert res == 421
+
     def test_dict_setdefault(self):
         def f():
             d = {}