Merge pull request #1511 from arjun23496/count_in_replace

Fixes #1495 - Added count argument for replacing contents in body

1 files changed, 7 insertions, 3 deletions

diff --git a/netlib/http/headers.py b/netlib/http/headers.py
index 36e5060c..131e8ce5 100644
--- a/netlib/http/headers.py
+++ b/netlib/http/headers.py
@@ -158,7 +158,7 @@ class Headers(multidict.MultiDict):
         else:
             return super(Headers, self).items()
 
-    def replace(self, pattern, repl, flags=0):
+    def replace(self, pattern, repl, flags=0, count=0):
         """
         Replaces a regular expression pattern with repl in each "name: value"
         header line.
@@ -172,10 +172,10 @@ class Headers(multidict.MultiDict):
             repl = strutils.escaped_str_to_bytes(repl)
         pattern = re.compile(pattern, flags)
         replacements = 0
-
+        flag_count = count > 0
         fields = []
         for name, value in self.fields:
-            line, n = pattern.subn(repl, name + b": " + value)
+            line, n = pattern.subn(repl, name + b": " + value, count=count)
             try:
                 name, value = line.split(b": ", 1)
             except ValueError:
@@ -184,6 +184,10 @@ class Headers(multidict.MultiDict):
                 pass
             else:
                 replacements += n
+                if flag_count:
+                    count -= n
+                    if count == 0:
+                        break
             fields.append((name, value))
         self.fields = tuple(fields)
         return replacements