replace() methods now decode and re-encode contents before substitution.

This commit is contained in:
Aldo Cortesi 2012-03-16 11:24:18 +13:00
parent d51b8cab0c
commit d138af7217
2 changed files with 41 additions and 9 deletions

View File

@ -176,8 +176,11 @@ class ODict:
def replace(self, pattern, repl, *args, **kwargs):
"""
Replaces a regular expression pattern with repl in both keys
and values. Returns the number of replacements made.
Replaces a regular expression pattern with repl in both keys and
values. Encoded content will be decoded before replacement, and
re-encoded afterwards.
Returns the number of replacements made.
"""
nlst, count = [], 0
for i in self.lst:
@ -475,9 +478,12 @@ class Request(HTTPMsg):
def replace(self, pattern, repl, *args, **kwargs):
"""
Replaces a regular expression pattern with repl in both the headers
and the body of the request. Returns the number of replacements
made.
and the body of the request. Encoded content will be decoded before
replacement, and re-encoded afterwards.
Returns the number of replacements made.
"""
with decoded(self):
self.content, c = re.subn(pattern, repl, self.content, *args, **kwargs)
self.path, pc = re.subn(pattern, repl, self.path, *args, **kwargs)
c += pc
@ -630,9 +636,12 @@ class Response(HTTPMsg):
def replace(self, pattern, repl, *args, **kwargs):
"""
Replaces a regular expression pattern with repl in both the headers
and the body of the response. Returns the number of replacements
made.
and the body of the response. Encoded content will be decoded
before replacement, and re-encoded afterwards.
Returns the number of replacements made.
"""
with decoded(self):
self.content, c = re.subn(pattern, repl, self.content, *args, **kwargs)
c += self.headers.replace(pattern, repl, *args, **kwargs)
return c
@ -753,6 +762,8 @@ class Error(controller.Msg):
Replaces a regular expression pattern with repl in both the headers
and the body of the request. Returns the number of replacements
made.
FIXME: Is replace useful on an Error object??
"""
self.msg, c = re.subn(pattern, repl, self.msg, *args, **kwargs)
return c
@ -1062,7 +1073,10 @@ class Flow:
def replace(self, pattern, repl, *args, **kwargs):
"""
Replaces a regular expression pattern with repl in all parts of the
flow . Returns the number of replacements made.
flow. Encoded content will be decoded before replacement, and
re-encoded afterwards.
Returns the number of replacements made.
"""
c = self.request.replace(pattern, repl, *args, **kwargs)
if self.response:

View File

@ -278,6 +278,24 @@ class uFlow(libpry.AutoTree):
f.replace("error", "bar")
assert f.error.msg == "bar"
def test_replace_encoded(self):
f = tutils.tflow_full()
f.request.content = "afoob"
f.request.encode("gzip")
f.response.content = "afoob"
f.response.encode("gzip")
f.replace("foo", "bar")
assert f.request.content != "abarb"
f.request.decode()
assert f.request.content == "abarb"
assert f.response.content != "abarb"
f.response.decode()
assert f.response.content == "abarb"
class uState(libpry.AutoTree):
def test_backup(self):