Allow custom authorization mechanisms.

docs/examples/custom_auth.py

@@ -0,0 +1,32 @@
+from __future__ import print_function
+from twisted.internet.task import react
+from twisted.web.iweb import IAgent
+from _utils import print_response
+from zope.interface import implementer
+
+import treq
+
+
+@implementer(IAgent)
+class CustomAuth(object):
+    """
+    I implement a custom authorization scheme.
+    """
+
+    def __init__(self, agent):
+        self._agent = agent
+
+    def request(self, method, uri, headers=None, bodyProducer=None):
+        print("Some authorization occurs here")
+        return self._agent.request(method, uri, headers, bodyProducer)
+
+
+def main(reactor, *args):
+    d = treq.get(
+        'http://httpbin.org/get',
+        auth=CustomAuth,
+    )
+    d.addCallback(print_response)
+    return d
+
+react(main, [])

docs/howto.rst

@@ -54,6 +54,18 @@ Full example: :download:`basic_auth.py <examples/basic_auth.py>`
 
 .. _RFC 2617: http://www.ietf.org/rfc/rfc2617.txt
 
+The ``auth`` keyword argument also supports custom authorization
+implementations.  ``auth`` should be a callable that accepts as its
+only argument the currently configured ``IAgent`` and should return an
+``IAgent`` that implements the desired authorization mechanism.
+
+.. literalinclude:: examples/custom_auth.py
+    :linenos:
+    :lines: 10-33
+
+Full example: :download:`basic_auth.py <examples/custom_auth.py>`
+
+
 Redirects
 ---------
 

treq/auth.py

@@ -37,5 +37,7 @@ def add_basic_auth(agent, username, password):
 def add_auth(agent, auth_config):
     if isinstance(auth_config, tuple):
         return add_basic_auth(agent, auth_config[0], auth_config[1])
+    elif callable(auth_config):
+        return auth_config(agent)
 
     raise UnknownAuthConfig(auth_config)

treq/test/test_auth.py

@@ -2,10 +2,13 @@
 
 from twisted.web.client import Agent
 from twisted.web.http_headers import Headers
+from twisted.web.iweb import IAgent
 
 from treq.test.util import TestCase
 from treq.auth import _RequestHeaderSettingAgent, add_auth, UnknownAuthConfig
 
+from zope.interface import implementer
+
 
 class RequestHeaderSettingAgentTests(TestCase):
     def setUp(self):
@@ -70,6 +73,22 @@ def test_add_basic_auth_huge(self):
             agent,
             Headers({b'authorization': [auth]}))
 
+    def test_auth_callable(self):
+        agent = mock.Mock()
+
+        @implementer(IAgent)
+        class AuthorizingAgent(object):
+
+            def __init__(self, agent):
+                self.wrapped_agent = agent
+
+            def request(self, method, uri, headers=None, bodyProducer=None):
+                """Not called by this test"""
+
+        wrapping_agent = add_auth(agent, AuthorizingAgent)
+        self.assertIsInstance(wrapping_agent, AuthorizingAgent)
+        self.assertIs(wrapping_agent.wrapped_agent, agent)
+
     def test_add_unknown_auth(self):
         agent = mock.Mock()
-        self.assertRaises(UnknownAuthConfig, add_auth, agent, mock.Mock())
+        self.assertRaises(UnknownAuthConfig, add_auth, agent, object())