Better way to handle complex object, i hope (refs #11)

Author: vimalloc

flask_jwt_extended/jwt_manager.py

@@ -8,7 +8,10 @@
 class JWTManager:
     def __init__(self, app=None):
         # Function that will be called to add custom user claims to a JWT.
-        self.user_claims_callback = lambda _: {}
+        self._user_claims_callback = lambda _: {}
+
+        # Function that will be called to return an identity from an object
+        self._user_identity_callback = lambda i: i
 
         # Function that will be called when an expired token is received
         self._expired_token_callback = lambda: (
@@ -90,7 +93,19 @@ def user_claims_loader(self, callback):
         Callback must be a function that takes only one argument, which is the
         identity of the JWT being created.
         """
-        self.user_claims_callback = callback
+        self._user_claims_callback = callback
+        return callback
+
+    def user_identity_loader(self, callback):
+        """
+        This sets the callback method for adding custom user claims to a JWT.
+
+        By default, no extra user claims will be added to the JWT.
+
+        Callback must be a function that takes only one argument, which is the
+        identity of the JWT being created.
+        """
+        self._user_identity_callback = callback
         return callback
 
     def expired_token_loader(self, callback):

flask_jwt_extended/utils.py

@@ -312,30 +312,28 @@ def create_refresh_token(identity):
     return refresh_token
 
 
-def create_access_token(identity, fresh=False, identity_lookup=None):
+def create_access_token(identity, fresh=False):
     """
     Creates a new access token
 
     :param identity: The identity of this token. This can be any data that is
                      json serializable. It can also be an object, in which case
-                     you can pass a function to identity_lookup which tells us
-                     how to get the identity out of this object. This is useful
-                     so you don't need to query disk twice, once for initially
-                     finding the identity in your login endpoint, and once for
-                     setting addition data in the JWT via the user_claims_loader
+                     you can use the user_identity_loader to define a function
+                     that will be called to pull a json serializable identity
+                     out of this object. This is useful so you don't need to
+                     query disk twice, once for initially finding the identity
+                     in your login endpoint, and once for setting addition data
+                     in the JWT via the user_claims_loader
     :param fresh: If this token should me markded as fresh, and can thus access
                   fresh_jwt_required protected endpoints. Defaults to False
-    :param identity_lookup: Function to generate a json serilizable identity
-                            from the identity object
     :return: A newly encoded JWT access token
     """
     # Token options
     secret = _get_secret_key()
     access_expire_delta = get_access_expires()
     algorithm = get_algorithm()
-    user_claims = current_app.jwt_manager.user_claims_callback(identity)
-    if identity_lookup:
-        identity = identity_lookup(identity)
+    user_claims = current_app.jwt_manager._user_claims_callback(identity)
+    identity = current_app.jwt_manager._user_identity_callback(identity)
 
     access_token = _encode_access_token(identity, secret, algorithm, access_expire_delta,
                                         fresh=fresh, user_claims=user_claims)

tests/test_jwt_encode_decode.py

@@ -336,7 +336,7 @@ def test_create_access_token_with_object(self):
         # Complex object to test building a JWT from. Normally if you are using
         # this functionality, this is something that would be retrieved from
         # disk somewhere (think sqlalchemy)
-        class TestObject:
+        class TestUser:
             def __init__(self, username, roles):
                 self.username = username
                 self.roles = roles
@@ -348,16 +348,19 @@ def __init__(self, username, roles):
         jwt = JWTManager(app)
 
         @jwt.user_claims_loader
-        def custom_claims(object):
+        def custom_claims(user):
             return {
-                'roles': object.roles
+                'roles': user.roles
             }
 
+        @jwt.user_identity_loader
+        def user_identity_lookup(user):
+            return user.username
+
         # Create the token using the complex object
         with app.test_request_context():
-            user = TestObject(username='foo', roles=['bar', 'baz'])
-            token = create_access_token(identity=user,
-                                        identity_lookup=lambda obj: obj.username)
+            user = TestUser(username='foo', roles=['bar', 'baz'])
+            token = create_access_token(identity=user)
 
             # Decode the token and make sure the values are set properly
             token_data = _decode_jwt(token, app.secret_key, app.config['JWT_ALGORITHM'])

tests/test_jwt_manager.py

@@ -32,7 +32,7 @@ def test_class_init(self):
     def test_default_user_claims_callback(self):
         identity = 'foobar'
         m = JWTManager(self.app)
-        assert m.user_claims_callback(identity) == {}
+        assert m._user_claims_callback(identity) == {}
 
     def test_default_expired_token_callback(self):
         with self.app.test_request_context():
@@ -88,7 +88,7 @@ def test_custom_user_claims_callback(self):
         def custom_user_claims(identity):
             return {'foo': 'bar'}
 
-        assert m.user_claims_callback(identity) == {'foo': 'bar'}
+        assert m._user_claims_callback(identity) == {'foo': 'bar'}
 
     def test_custom_expired_token_callback(self):
         with self.app.test_request_context():