From ee909e265b40743479ca2f3e8b518f76c91f83f8 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Sat, 28 Apr 2012 22:54:45 +1200
Subject: [PATCH] Add a shortcut for setting Location header.

---
 libpathod/rparse.py | 21 +++++++++++++++++++++
 test/test_rparse.py |  7 ++++---
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/libpathod/rparse.py b/libpathod/rparse.py
index 4f7cc12d2..b47681cc2 100644
--- a/libpathod/rparse.py
+++ b/libpathod/rparse.py
@@ -211,6 +211,26 @@ class ShortcutContentType:
         return e.setParseAction(lambda x: klass(*x))
 
 
+
+class ShortcutLocation:
+    def __init__(self, value):
+        self.value = value
+
+    def mod_response(self, settings, r):
+        r.headers.append(
+            (
+                LiteralGenerator("Location"),
+                self.value.get_generator(settings)
+            )
+        )
+
+    @classmethod
+    def expr(klass):
+        e = pp.Literal("l").suppress()
+        e = e + Value
+        return e.setParseAction(lambda x: klass(*x))
+
+
 class Body:
     def __init__(self, value):
         self.value = value
@@ -332,6 +352,7 @@ class Response:
         DisconnectBefore,
         DisconnectRandom,
         ShortcutContentType,
+        ShortcutLocation,
     )
     version = "HTTP/1.1"
     code = 200
diff --git a/test/test_rparse.py b/test/test_rparse.py
index 5a3cb53cd..c854b0c6f 100644
--- a/test/test_rparse.py
+++ b/test/test_rparse.py
@@ -177,9 +177,10 @@ class uDisconnects(libpry.AutoTree):
         assert isinstance(v, rparse.DisconnectRandom)
 
 
-class uShortcutContentType(libpry.AutoTree):
+class uShortcuts(libpry.AutoTree):
     def test_parse(self):
-        assert len(rparse.parse({}, "400:c'foo'").headers) == 1
+        assert rparse.parse({}, "400:c'foo'").headers[0][0][:] == "Content-Type"
+        assert rparse.parse({}, "400:l'foo'").headers[0][0][:] == "Location"
 
 
 class uPauses(libpry.AutoTree):
@@ -312,5 +313,5 @@ tests = [
     uDisconnects(),
     uMisc(),
     uparse(),
-    uShortcutContentType()
+    uShortcuts()
 ]