Merge branch 'better-replace'

3 files changed, 66 insertions, 1 deletions

diff --git a/netlib/http/headers.py b/netlib/http/headers.py
index bcb828da..72739f90 100644
--- a/netlib/http/headers.py
+++ b/netlib/http/headers.py
@@ -6,6 +6,8 @@ See also: http://lucumr.pocoo.org/2013/7/2/the-updated-guide-to-unicode/
 """
 from __future__ import absolute_import, print_function, division
 
+import re
+
 try:
     from collections.abc import MutableMapping
 except ImportError:  # pragma: no cover
@@ -198,4 +200,31 @@ class Headers(MutableMapping, Serializable):
 
     @classmethod
     def from_state(cls, state):
-        return cls([list(field) for field in state])
\ No newline at end of file
+        return cls([list(field) for field in state])
+
+    @_always_byte_args
+    def replace(self, pattern, repl, flags=0):
+        """
+        Replaces a regular expression pattern with repl in each "name: value"
+        header line.
+
+        Returns:
+            The number of replacements made.
+        """
+        pattern = re.compile(pattern, flags)
+        replacements = 0
+
+        fields = []
+        for name, value in self.fields:
+            line, n = pattern.subn(repl, name + b": " + value)
+            try:
+                name, value = line.split(b": ", 1)
+            except ValueError:
+                # We get a ValueError if the replacement removed the ": "
+                # There's not much we can do about this, so we just keep the header as-is.
+                pass
+            else:
+                replacements += n
+            fields.append([name, value])
+        self.fields = fields
+        return replacements
diff --git a/netlib/http/message.py b/netlib/http/message.py
index b265ac4f..da9681a0 100644
--- a/netlib/http/message.py
+++ b/netlib/http/message.py
@@ -175,6 +175,25 @@ class Message(utils.Serializable):
         self.headers["content-encoding"] = e
         return True
 
+    def replace(self, pattern, repl, flags=0):
+        """
+        Replaces a regular expression pattern with repl in both the headers
+        and the body of the message. Encoded body will be decoded
+        before replacement, and re-encoded afterwards.
+
+        Returns:
+            The number of replacements made.
+        """
+        # TODO: Proper distinction between text and bytes.
+        replacements = 0
+        if self.content:
+            with decoded(self):
+                self.content, replacements = utils.safe_subn(
+                    pattern, repl, self.content, flags=flags
+                )
+        replacements += self.headers.replace(pattern, repl, flags)
+        return replacements
+
     # Legacy
 
     @property
diff --git a/netlib/http/request.py b/netlib/http/request.py
index db4901b7..692ba30f 100644
--- a/netlib/http/request.py
+++ b/netlib/http/request.py
@@ -54,6 +54,23 @@ class Request(Message):
             self.method, hostport, path
         )
 
+    def replace(self, pattern, repl, flags=0):
+        """
+            Replaces a regular expression pattern with repl in the headers, the
+            request path and the body of the request. Encoded content will be
+            decoded before replacement, and re-encoded afterwards.
+
+            Returns:
+                The number of replacements made.
+        """
+        # TODO: Proper distinction between text and bytes.
+        c = super(Request, self).replace(pattern, repl, flags)
+        self.path, pc = utils.safe_subn(
+            pattern, repl, self.path, flags=flags
+        )
+        c += pc
+        return c
+
     @property
     def first_line_format(self):
         """