Use bytes or unicode to get args

src/klein/app.py

@@ -24,7 +24,7 @@ def iscoroutine(*args, **kwargs):
 
 from twisted.web.iweb import IRenderable
 from twisted.web.template import renderElement
-from twisted.web.server import Site, Request
+from twisted.web.server import Request
 from twisted.internet import reactor, endpoints
 
 try:
@@ -35,7 +35,7 @@ def ensureDeferred(*args, **kwagrs):
 
 from zope.interface import implementer
 
-from klein.resource import KleinResource
+from klein.resource import KleinResource, KleinSite
 from klein.interfaces import IKleinRequest
 
 __all__ = ['Klein', 'run', 'route', 'resource']
@@ -367,7 +367,7 @@ def run(self, host=None, port=None, logFile=None,
                                                                        host)
 
         endpoint = endpoints.serverFromString(reactor, endpoint_description)
-        endpoint.listen(Site(self.resource()))
+        endpoint.listen(KleinSite(self.resource()))
         reactor.run()
 
 

src/klein/resource.py

@@ -5,7 +5,7 @@
 from twisted.internet import defer
 from twisted.python import log, failure
 from twisted.python.compat import unicode, intToBytes
-from twisted.web import server
+from twisted.web import http, server
 from twisted.web.iweb import IRenderable
 from twisted.web.resource import Resource, IResource, getChildForRequest
 from twisted.web.server import NOT_DONE_YET
@@ -280,3 +280,39 @@ def write_response(r):
         d.addCallback(write_response).addErrback(log.err, _why="Unhandled Error writing response")
 
         return server.NOT_DONE_YET
+
+
+
+class KleinHTTPRequest(server.Request):
+
+    def getArg(self, key):
+        """
+        Get a single arg value.
+
+        @raises KeyError: If key doesn't exist
+        @raises ValueError: If there is more than 1 value
+        @return: L{list} of L{bytes}
+        """
+        key = ensure_utf8_bytes(key)
+        value = self.args[key]
+        if len(value) != 1:
+            raise ValueError('Too many values for: {0}'.format(key))
+        return value[0]
+
+    def getArgs(self, key):
+        """
+        Get the list of values for a key.
+
+        @return: L{list} of L{bytes}
+        """
+        key = ensure_utf8_bytes(key)
+        return self.args.get(key, [])
+
+
+
+class KleinSite(server.Site):
+    def buildProtocol(self, addr):
+        channel = http.HTTPFactory.buildProtocol(self, addr)
+        channel.requestFactory = KleinHTTPRequest
+        channel.site = self
+        return channel

src/klein/test/test_app.py

@@ -246,12 +246,12 @@ def foo(request):
 
 
     @patch('klein.app.KleinResource')
-    @patch('klein.app.Site')
+    @patch('klein.app.KleinSite')
     @patch('klein.app.log')
     @patch('klein.app.reactor')
     def test_run(self, reactor, mock_log, mock_site, mock_kr):
         """
-        L{Klein.run} configures a L{KleinResource} and a L{Site}
+        L{Klein.run} configures a L{KleinResource} and a L{KleinSite}
         listening on the specified interface and port, and logs
         to stdout.
         """
@@ -270,7 +270,7 @@ def test_run(self, reactor, mock_log, mock_site, mock_kr):
 
 
     @patch('klein.app.KleinResource')
-    @patch('klein.app.Site')
+    @patch('klein.app.KleinSite')
     @patch('klein.app.log')
     @patch('klein.app.reactor')
     def test_runWithLogFile(self, reactor, mock_log, mock_site, mock_kr):

src/klein/test/test_request.py

@@ -0,0 +1,61 @@
+from __future__ import absolute_import
+
+from klein import Klein
+from klein.resource import KleinHTTPRequest, KleinSite
+from klein.test.util import TestCase
+from twisted.web.test.requesthelper import DummyChannel
+
+class BytesUnicodeTest(TestCase):
+
+    def setUp(self):
+        self.encoding = 'utf-8'
+        app = Klein()
+        site = KleinSite(app.resource())
+        channel = site.buildProtocol('127.0.0.1')
+        self.request = channel.requestFactory(DummyChannel(), None)
+        self.request.args = {}
+
+    def test_getArg(self):
+        str_key = 'test'
+        bytes_key = str_key.encode(self.encoding)
+        value = b'hello world'
+
+        self.request.args[bytes_key] = [value]
+        self.assertEquals(self.request.getArg(str_key), value)
+        self.assertEquals(self.request.getArg(bytes_key), value)
+
+    def test_getArg_not_1(self):
+        """
+        Raise exception if there are more or less values than 1
+        """
+        str_key = 'test'
+        bytes_key = str_key.encode(self.encoding)
+        values = []
+
+        self.request.args[bytes_key] = values
+        self.assertRaises(ValueError, self.request.getArg, str_key)
+        self.assertRaises(ValueError, self.request.getArg, bytes_key)
+
+        values.extend([b'hello', b'world'])
+        self.assertRaises(ValueError, self.request.getArg, str_key)
+        self.assertRaises(ValueError, self.request.getArg, bytes_key)
+
+    def test_getArgs(self):
+        str_key = 'test'
+        bytes_key = str_key.encode(self.encoding)
+        values = [b'hello world', b'hey earth']
+        self.request.args[bytes_key] = values
+
+        self.assertEquals(len(self.request.getArgs(str_key)), len(values))
+        self.assertEquals(self.request.getArgs(str_key), values)
+        self.assertEquals(self.request.getArgs(bytes_key), values)
+
+    def test_getArgs_no_key(self):
+        """
+        By default, an empty list is returned if a key doesn't exist
+        """
+        str_key = 'test'
+        bytes_key = str_key.encode(self.encoding)
+
+        self.assertEquals(self.request.getArgs(str_key), [])
+        self.assertEquals(self.request.getArgs(bytes_key), [])