[pypy-commit] pypy default: catch all the needed memoryerror conditions in str.replace

[alex_gaynor]

Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: 
Changeset: r50818:072b5317cc0f
Date: 2011-12-21 21:14 -0600
http://bitbucket.org/pypy/pypy/changeset/072b5317cc0f/

Log:	catch all the needed memoryerror conditions in str.replace

diff --git a/pypy/objspace/std/stringobject.py b/pypy/objspace/std/stringobject.py
--- a/pypy/objspace/std/stringobject.py
+++ b/pypy/objspace/std/stringobject.py
@@ -510,6 +510,16 @@
 
     return space.wrap(res)
 
+def _replace_overflow_check(space, builder, new_piece):
+    # Checks if adding new_piece chars to the builder would overflow, and
+    # converts into an OverflowError.
+    try:
+        ovfcheck(builder.getlength() + new_piece)
+    except OverflowError:
+        raise OperationError(space.w_OverflowError,
+            space.wrap("replace string is too long")
+        )
+
 def _string_replace(space, input, sub, by, maxsplit):
     if maxsplit == 0:
         return space.wrap(input)
@@ -519,7 +529,7 @@
         if maxsplit > 0 and maxsplit < upper + 2:
             upper = maxsplit - 1
             assert upper >= 0
-        
+
         try:
             result_size = ovfcheck(upper * len(by))
             result_size = ovfcheck(result_size + upper)
@@ -548,14 +558,18 @@
             if next < 0:
                 break
             if not first:
+                _replace_overflow_check(space, builder, len(by))
                 builder.append(by)
             first = False
+            _replace_overflow_check(space, builder, next - start)
             builder.append_slice(input, start, next)
             start = next + sublen
             maxsplit -= 1   # NB. if it's already < 0, it stays < 0
 
         if not first:
+            _replace_overflow_check(space, builder, len(by))
             builder.append(by)
+        _replace_overflow_check(space, builder, len(input) - start)
         builder.append_slice(input, start, len(input))
 
     return space.wrap(builder.build())