Add a decoded context manager.

This simplifies a common chore when modifying traffic - decoding the object,
modifying it, then re-encoding it with the same encoding afterwards. You can
now simply say:

with flow.decoded(request):
    request.content = "bar"
This commit is contained in:
Aldo Cortesi 2012-03-16 11:12:56 +13:00
parent 8d662e6636
commit d51b8cab0c
2 changed files with 52 additions and 0 deletions

View File

@ -199,6 +199,34 @@ class ODictCaseless(ODict):
return s.lower() return s.lower()
class decoded(object):
"""
A context manager that decodes a request, response or error, and then
re-encodes it with the same encoding after execution of the block.
Example:
with decoded(request):
request.content = request.content.replace("foo", "bar")
"""
def __init__(self, o):
self.o = o
ce = o.headers["content-encoding"]
if ce and ce[0] in encoding.ENCODINGS:
self.ce = ce[0]
else:
self.ce = None
def __enter__(self):
if self.ce:
self.o.decode()
def __exit__(self, type, value, tb):
if self.ce:
self.o.encode(self.ce)
class HTTPMsg(controller.Msg): class HTTPMsg(controller.Msg):
def decode(self): def decode(self):
""" """

View File

@ -996,6 +996,29 @@ class uODictCaseless(libpry.AutoTree):
assert len(self.od) == 1 assert len(self.od) == 1
class udecoded(libpry.AutoTree):
def test_del(self):
r = tutils.treq()
assert r.content == "content"
assert not r.headers["content-encoding"]
r.encode("gzip")
assert r.headers["content-encoding"]
assert r.content != "content"
with flow.decoded(r):
assert not r.headers["content-encoding"]
assert r.content == "content"
assert r.headers["content-encoding"]
assert r.content != "content"
with flow.decoded(r):
r.content = "foo"
assert r.content != "foo"
r.decode()
assert r.content == "foo"
tests = [ tests = [
uStickyCookieState(), uStickyCookieState(),
@ -1012,4 +1035,5 @@ tests = [
uClientConnect(), uClientConnect(),
uODict(), uODict(),
uODictCaseless(), uODictCaseless(),
udecoded()
] ]