[pypy-commit] pypy default: (ronan, fijal, esoda) strip for rpython

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r68257:973ffa4ffa68
Date: 2013-11-20 11:59 +0100
http://bitbucket.org/pypy/pypy/changeset/973ffa4ffa68/

Log:	(ronan, fijal, esoda) strip for rpython

diff --git a/rpython/annotator/unaryop.py b/rpython/annotator/unaryop.py
--- a/rpython/annotator/unaryop.py
+++ b/rpython/annotator/unaryop.py
@@ -460,13 +460,13 @@
         check_negative_slice(start, end, "count")
         return SomeInteger(nonneg=True)
 
-    def method_strip(str, chr):
+    def method_strip(str, chr=None):
         return str.basestringclass(no_nul=str.no_nul)
 
-    def method_lstrip(str, chr):
+    def method_lstrip(str, chr=None):
         return str.basestringclass(no_nul=str.no_nul)
 
-    def method_rstrip(str, chr):
+    def method_rstrip(str, chr=None):
         return str.basestringclass(no_nul=str.no_nul)
 
     def method_join(str, s_list):
diff --git a/rpython/rtyper/lltypesystem/rstr.py b/rpython/rtyper/lltypesystem/rstr.py
--- a/rpython/rtyper/lltypesystem/rstr.py
+++ b/rpython/rtyper/lltypesystem/rstr.py
@@ -9,6 +9,7 @@
 from rpython.rlib.rarithmetic import ovfcheck
 from rpython.rtyper.error import TyperError
 from rpython.rtyper.lltypesystem import ll_str, llmemory
+from rpython.rtyper.lltypesystem.lloperation import llop
 from rpython.rtyper.lltypesystem.lltype import (GcStruct, Signed, Array, Char,
     UniChar, Ptr, malloc, Bool, Void, GcArray, nullptr, cast_primitive,
     typeOf, staticAdtMethod, GcForwardReference)
@@ -402,6 +403,46 @@
         return result
 
     @jit.elidable
+    def ll_strip_default(s, left, right):
+        s_len = len(s.chars)
+        if s_len == 0:
+            return s.empty()
+        lpos = 0
+        rpos = s_len - 1
+        if left:
+            while lpos < rpos and s.chars[lpos].isspace():
+                lpos += 1
+        if right:
+            while lpos < rpos + 1 and s.chars[rpos].isspace():
+                rpos -= 1
+        if rpos < lpos:
+            return s.empty()
+        r_len = rpos - lpos + 1
+        result = s.malloc(r_len)
+        s.copy_contents(s, result, lpos, 0, r_len)
+        return result
+
+    @jit.elidable
+    def ll_strip_multiple(s, s2, left, right):
+        s_len = len(s.chars)
+        if s_len == 0:
+            return s.empty()
+        lpos = 0
+        rpos = s_len - 1
+        if left:
+            while lpos < rpos and LLHelpers.ll_contains(s2, s.chars[lpos]):
+                lpos += 1
+        if right:
+            while lpos < rpos + 1 and LLHelpers.ll_contains(s2, s.chars[rpos]):
+                rpos -= 1
+        if rpos < lpos:
+            return s.empty()
+        r_len = rpos - lpos + 1
+        result = s.malloc(r_len)
+        s.copy_contents(s, result, lpos, 0, r_len)
+        return result
+
+    @jit.elidable
     def ll_upper(s):
         s_chars = s.chars
         s_len = len(s_chars)
diff --git a/rpython/rtyper/rstr.py b/rpython/rtyper/rstr.py
--- a/rpython/rtyper/rstr.py
+++ b/rpython/rtyper/rstr.py
@@ -231,11 +231,22 @@
     def rtype_method_strip(self, hop, left=True, right=True):
         rstr = hop.args_r[0].repr
         v_str = hop.inputarg(rstr.repr, arg=0)
-        v_char = hop.inputarg(rstr.char_repr, arg=1)
-        v_left = hop.inputconst(Bool, left)
-        v_right = hop.inputconst(Bool, right)
+        args_v = [v_str]
+        if len(hop.args_s) == 2:
+            if isinstance(hop.args_s[1], annmodel.SomeString):
+                v_stripstr = hop.inputarg(rstr.repr, arg=1)
+                args_v.append(v_stripstr)
+                func = self.ll.ll_strip_multiple
+            else:
+                v_char = hop.inputarg(rstr.char_repr, arg=1)
+                args_v.append(v_char)
+                func = self.ll.ll_strip
+        else:
+            func = self.ll.ll_strip_default
+        args_v.append(hop.inputconst(Bool, left))
+        args_v.append(hop.inputconst(Bool, right))
         hop.exception_is_here()
-        return hop.gendirectcall(self.ll.ll_strip, v_str, v_char, v_left, v_right)
+        return hop.gendirectcall(func, *args_v)
 
     def rtype_method_lstrip(self, hop):
         return self.rtype_method_strip(hop, left=True, right=False)
diff --git a/rpython/rtyper/test/test_rstr.py b/rpython/rtyper/test/test_rstr.py
--- a/rpython/rtyper/test/test_rstr.py
+++ b/rpython/rtyper/test/test_rstr.py
@@ -9,6 +9,7 @@
 from rpython.rtyper.rstr import AbstractLLHelpers
 from rpython.rtyper.rtyper import TyperError
 from rpython.rtyper.test.tool import BaseRtypingTest
+from rpython.rtyper.annlowlevel import llstr, hlstr
 
 
 def test_parse_fmt():
@@ -457,6 +458,29 @@
         res = self.interpret(left2, [])
         assert self.ll_to_string(res) == const('a')
 
+    def test_strip_multiple_chars(self):
+        const = self.const
+        def both():
+            return const('!ab!').strip(const('!a'))
+        def left():
+            return const('!+ab!').lstrip(const('!+'))
+        def right():
+            return const('!ab!+').rstrip(const('!+'))
+        def empty():
+            return const(' \t\t   ').strip('\t ')
+        def left2():
+            return const('a  ').strip(' \t')
+        res = self.interpret(both, [])
+        assert self.ll_to_string(res) == const('b')
+        res = self.interpret(left, [])
+        assert self.ll_to_string(res) == const('ab!')
+        res = self.interpret(right, [])
+        assert self.ll_to_string(res) == const('!ab')
+        res = self.interpret(empty, [])
+        assert self.ll_to_string(res) == const('')
+        res = self.interpret(left2, [])
+        assert self.ll_to_string(res) == const('a')
+
     def test_upper(self):
         const = self.const
         constchar = self.constchar
@@ -1143,3 +1167,16 @@
         self.interpret(f, [array, 4])
         assert list(array) == list('abc'*4)
         lltype.free(array, flavor='raw')
+
+    def test_strip_no_arg(self):
+        strings = ["  xyz  ", "", "\t\vx"]
+
+        def f(i):
+            return strings[i].strip()
+
+        res = self.interpret(f, [0])
+        assert hlstr(res) == "xyz"
+        res = self.interpret(f, [1])
+        assert hlstr(res) == ""
+        res = self.interpret(f, [2])
+        assert hlstr(res) == "x"