mymemreplace(), strop_replace(): Add support for optional 4th argument

maxsplit which is implemented in string.py but wasn't here.  The
reference manual doesn't define what happens when maxsplit is negative
or larger than the number of occurrences, but in either case, I
implemented this as all get replaced.  Default value is zero which
replaces all occurrences.

1 files changed, 12 insertions, 4 deletions

diff --git a/Modules/stropmodule.c b/Modules/stropmodule.c
index d52ac7f..9fc4b9b 100644
--- a/Modules/stropmodule.c
+++ b/Modules/stropmodule.c
@@ -877,13 +877,14 @@ static int mymemcnt(mem, len, pat, pat_len)
        the new string allocated locally, or
        NULL if an error occurred.
 */
-static char *mymemreplace(str, len, pat, pat_len, sub, sub_len, out_len)
+static char *mymemreplace(str, len, pat, pat_len, sub, sub_len, count, out_len)
 	char *str;
 	int len;     /* input string  */
 	char *pat;
 	int pat_len; /* pattern string to find */
 	char *sub;
 	int sub_len; /* substitution string */
+	int count;   /* number of replacements, 0 == all */
 	int *out_len;
 
 {
@@ -896,6 +897,8 @@ static char *mymemreplace(str, len, pat, pat_len, sub, sub_len, out_len)
 
 	/* find length of output string */
 	nfound = mymemcnt(str, len, pat, pat_len);
+	if (count > 0)
+		nfound = nfound > count ? count : nfound;
 	if (nfound == 0)
 		goto return_same;
 	new_len = len + nfound*(sub_len - pat_len);
@@ -921,6 +924,9 @@ static char *mymemreplace(str, len, pat, pat_len, sub, sub_len, out_len)
 		new_s += offset; /* move new_s to dest for sub string */
 		memcpy(new_s, sub, sub_len); /* copy substring into new_s */
 		new_s += sub_len; /* offset new_s past sub string */
+
+		/* break when we've done count replacements */
+		if (--count == 0) break;
 	}
 	/* copy any remaining values into output string */
 	if (len > 0)
@@ -940,12 +946,14 @@ strop_replace(self, args)
 {
 	char *str, *pat,*sub,*new_s;
 	int len,pat_len,sub_len,out_len;
+	int count = 0;
 	PyObject *new;
 
-	if (!PyArg_ParseTuple(args, "s#s#s#",
-			      &str, &len, &pat, &pat_len, &sub, &sub_len))
+	if (!PyArg_ParseTuple(args, "s#s#s#|i",
+			      &str, &len, &pat, &pat_len, &sub, &sub_len,
+			      &count))
 		return NULL;
-	new_s = mymemreplace(str,len,pat,pat_len,sub,sub_len,&out_len);
+	new_s = mymemreplace(str,len,pat,pat_len,sub,sub_len,count,&out_len);
 	if (new_s == NULL) {
 		PyErr_NoMemory();
 		return NULL;