Added replace for request and headers - Fixed missing replace option

2 files changed, 6 insertions, 5 deletions

diff --git a/netlib/http/headers.py b/netlib/http/headers.py
index 36e5060c..60a36515 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.
@@ -175,7 +175,7 @@ class Headers(multidict.MultiDict):
 
         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:
@@ -183,6 +183,7 @@ class Headers(multidict.MultiDict):
                 # There's not much we can do about this, so we just keep the header as-is.
                 pass
             else:
+                count -= n
                 replacements += n
             fields.append((name, value))
         self.fields = tuple(fields)
diff --git a/netlib/http/request.py b/netlib/http/request.py
index 666a5869..e0aaa8a9 100644
--- a/netlib/http/request.py
+++ b/netlib/http/request.py
@@ -80,7 +80,7 @@ class Request(message.Message):
             self.method, hostport, path
         )
 
-    def replace(self, pattern, repl, flags=0):
+    def replace(self, pattern, repl, flags=0, count=0):
         """
             Replaces a regular expression pattern with repl in the headers, the
             request path and the body of the request. Encoded content will be
@@ -94,9 +94,9 @@ class Request(message.Message):
         if isinstance(repl, six.text_type):
             repl = strutils.escaped_str_to_bytes(repl)
 
-        c = super(Request, self).replace(pattern, repl, flags)
+        c = super(Request, self).replace(pattern, repl, flags, count)
         self.path, pc = re.subn(
-            pattern, repl, self.data.path, flags=flags
+            pattern, repl, self.data.path, flags=flags, count=count
         )
         c += pc
         return c