[pypy-commit] pypy numpy-back-to-applevel: merge

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-back-to-applevel
Changeset: r51614:da00788e9d17
Date: 2012-01-22 00:10 +0200
http://bitbucket.org/pypy/pypy/changeset/da00788e9d17/

Log:	merge

diff --git a/pypy/annotation/unaryop.py b/pypy/annotation/unaryop.py
--- a/pypy/annotation/unaryop.py
+++ b/pypy/annotation/unaryop.py
@@ -440,6 +440,12 @@
     def method_popitem(dct):
         return dct.getanyitem('items')
 
+    def method_pop(dct, s_key, s_dfl=None):
+        dct.dictdef.generalize_key(s_key)
+        if s_dfl is not None:
+            dct.dictdef.generalize_value(s_dfl)
+        return dct.dictdef.read_value()
+
     def _can_only_throw(dic, *ignore):
         if dic1.dictdef.dictkey.custom_eq_hash:
             return None    # r_dict: can throw anything
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
@@ -323,6 +323,16 @@
         hop.exception_is_here()
         return hop.gendirectcall(ll_popitem, cTUPLE, v_dict)
 
+    def rtype_method_pop(self, hop):
+        if hop.nb_args == 2:
+            v_args = hop.inputargs(self, self.key_repr)
+            target = ll_pop
+        elif hop.nb_args == 3:
+            v_args = hop.inputargs(self, self.key_repr, self.value_repr)
+            target = ll_pop_default
+        hop.exception_is_here()
+        return hop.gendirectcall(target, *v_args)
+
 class __extend__(pairtype(DictRepr, rmodel.Repr)):
 
     def rtype_getitem((r_dict, r_key), hop):
@@ -874,3 +884,18 @@
     r.item1 = recast(ELEM.TO.item1, entry.value)
     _ll_dict_del(dic, i)
     return r
+
+def ll_pop(dic, key):
+    i = ll_dict_lookup(dic, key, dic.keyhash(key))
+    if not i & HIGHEST_BIT:
+        value = ll_get_value(dic, i)
+        _ll_dict_del(dic, i)
+        return value
+    else:
+        raise KeyError
+
+def ll_pop_default(dic, key, dfl):
+    try:
+        return ll_pop(dic, key)
+    except KeyError:
+        return dfl
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
@@ -160,6 +160,16 @@
         hop.exception_is_here()
         return hop.gendirectcall(ll_popitem, cTUPLE, v_dict)
 
+    def rtype_method_pop(self, hop):
+        if hop.nb_args == 2:
+            v_args = hop.inputargs(self, self.key_repr)
+            target = ll_pop
+        elif hop.nb_args == 3:
+            v_args = hop.inputargs(self, self.key_repr, self.value_repr)
+            target = ll_pop_default
+        hop.exception_is_here()
+        return hop.gendirectcall(target, *v_args)
+
     def __get_func(self, interp, r_func, fn, TYPE):
         if isinstance(r_func, MethodOfFrozenPBCRepr):
             obj = r_func.r_im_self.convert_const(fn.im_self)
@@ -370,6 +380,20 @@
         return res
     raise KeyError
 
+def ll_pop(d, key):
+    if d.ll_contains(key):
+        value = d.ll_get(key)
+        d.ll_remove(key)
+        return value
+    else:
+        raise KeyError
+
+def ll_pop_default(d, key, dfl):
+    try:
+        return ll_pop(d, key)
+    except KeyError:
+        return dfl
+
 # ____________________________________________________________
 #
 #  Iteration.
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
@@ -622,6 +622,26 @@
         res = self.interpret(func, [])
         assert res in [5263, 6352]
 
+    def test_dict_pop(self):
+        def f(n, default):
+            d = {}
+            d[2] = 3
+            d[4] = 5
+            if default == -1:
+                try:
+                    x = d.pop(n)
+                except KeyError:
+                    x = -1
+            else:
+                x = d.pop(n, default)
+            return x * 10 + len(d)
+        res = self.interpret(f, [2, -1])
+        assert res == 31
+        res = self.interpret(f, [3, -1])
+        assert res == -8
+        res = self.interpret(f, [2, 5])
+        assert res == 31
+
 class TestLLtype(BaseTestRdict, LLRtypeMixin):
     def test_dict_but_not_with_char_keys(self):
         def func(i):