Preserve types of functions decorated with jwt_required and loaders

docs/conf.py

@@ -366,4 +366,6 @@
     ("py:class", "flask.app.Flask"),
     ("py:class", "datetime.timedelta"),
     ("py:class", "flask.wrappers.Response"),
+    ("py:class", "flask_jwt_extended.jwt_manager._F"),
+    ("py:class", "flask_jwt_extended.view_decorators._F"),
 ]

flask_jwt_extended/jwt_manager.py

@@ -2,6 +2,7 @@
 from typing import Any
 from typing import Callable
 from typing import Optional
+from typing import TypeVar
 
 import jwt
 from flask import Flask
@@ -45,6 +46,8 @@
 from flask_jwt_extended.typing import Fresh
 from flask_jwt_extended.utils import current_user_context_processor
 
+_F = TypeVar("_F", bound=Callable[..., Any])
+
 
 class JWTManager(object):
     """
@@ -231,7 +234,7 @@ def _set_default_configuration_options(app: Flask) -> None:
         app.config.setdefault("JWT_VERIFY_SUB", True)
         app.config.setdefault("JWT_ENCODE_NBF", True)
 
-    def additional_claims_loader(self, callback: Callable) -> Callable:
+    def additional_claims_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function used to add additional claims
         when creating a JWT. The claims returned by this function will be merged
@@ -248,7 +251,7 @@ def additional_claims_loader(self, callback: Callable) -> Callable:
         self._user_claims_callback = callback
         return callback
 
-    def additional_headers_loader(self, callback: Callable) -> Callable:
+    def additional_headers_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function used to add additional headers
         when creating a JWT. The headers returned by this function will be merged
@@ -265,7 +268,7 @@ def additional_headers_loader(self, callback: Callable) -> Callable:
         self._jwt_additional_header_callback = callback
         return callback
 
-    def decode_key_loader(self, callback: Callable) -> Callable:
+    def decode_key_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function for dynamically setting the JWT
         decode key based on the **UNVERIFIED** contents of the token. Think
@@ -286,7 +289,7 @@ def decode_key_loader(self, callback: Callable) -> Callable:
         self._decode_key_callback = callback
         return callback
 
-    def encode_key_loader(self, callback: Callable) -> Callable:
+    def encode_key_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function for dynamically setting the JWT
         encode key based on the tokens identity. Think carefully before using this
@@ -302,7 +305,7 @@ def encode_key_loader(self, callback: Callable) -> Callable:
         self._encode_key_callback = callback
         return callback
 
-    def expired_token_loader(self, callback: Callable) -> Callable:
+    def expired_token_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function for returning a custom
         response when an expired JWT is encountered.
@@ -318,7 +321,7 @@ def expired_token_loader(self, callback: Callable) -> Callable:
         self._expired_token_callback = callback
         return callback
 
-    def invalid_token_loader(self, callback: Callable) -> Callable:
+    def invalid_token_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function for returning a custom
         response when an invalid JWT is encountered.
@@ -335,7 +338,7 @@ def invalid_token_loader(self, callback: Callable) -> Callable:
         self._invalid_token_callback = callback
         return callback
 
-    def needs_fresh_token_loader(self, callback: Callable) -> Callable:
+    def needs_fresh_token_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function for returning a custom
         response when a valid and non-fresh token is used on an endpoint
@@ -352,7 +355,7 @@ def needs_fresh_token_loader(self, callback: Callable) -> Callable:
         self._needs_fresh_token_callback = callback
         return callback
 
-    def revoked_token_loader(self, callback: Callable) -> Callable:
+    def revoked_token_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function for returning a custom
         response when a revoked token is encountered.
@@ -368,7 +371,7 @@ def revoked_token_loader(self, callback: Callable) -> Callable:
         self._revoked_token_callback = callback
         return callback
 
-    def token_in_blocklist_loader(self, callback: Callable) -> Callable:
+    def token_in_blocklist_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function used to check if a JWT has
         been revoked.
@@ -385,7 +388,7 @@ def token_in_blocklist_loader(self, callback: Callable) -> Callable:
         self._token_in_blocklist_callback = callback
         return callback
 
-    def token_verification_failed_loader(self, callback: Callable) -> Callable:
+    def token_verification_failed_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function used to return a custom
         response when the claims verification check fails.
@@ -401,7 +404,7 @@ def token_verification_failed_loader(self, callback: Callable) -> Callable:
         self._token_verification_failed_callback = callback
         return callback
 
-    def token_verification_loader(self, callback: Callable) -> Callable:
+    def token_verification_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function used for custom verification
         of a valid JWT.
@@ -418,7 +421,7 @@ def token_verification_loader(self, callback: Callable) -> Callable:
         self._token_verification_callback = callback
         return callback
 
-    def unauthorized_loader(self, callback: Callable) -> Callable:
+    def unauthorized_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function used to return a custom
         response when no JWT is present.
@@ -432,7 +435,7 @@ def unauthorized_loader(self, callback: Callable) -> Callable:
         self._unauthorized_callback = callback
         return callback
 
-    def user_identity_loader(self, callback: Callable) -> Callable:
+    def user_identity_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function used to convert an identity to
         a string when creating JWTs. This is useful for using objects (such as
@@ -447,7 +450,7 @@ def user_identity_loader(self, callback: Callable) -> Callable:
         self._user_identity_callback = callback
         return callback
 
-    def user_lookup_loader(self, callback: Callable) -> Callable:
+    def user_lookup_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function used to convert a JWT into
         a python object that can be used in a protected endpoint. This is useful
@@ -472,7 +475,7 @@ def user_lookup_loader(self, callback: Callable) -> Callable:
         self._user_lookup_callback = callback
         return callback
 
-    def user_lookup_error_loader(self, callback: Callable) -> Callable:
+    def user_lookup_error_loader(self, callback: _F) -> _F:
         """
         This decorator sets the callback function used to return a custom
         response when loading a user via

flask_jwt_extended/view_decorators.py

@@ -3,9 +3,12 @@
 from functools import wraps
 from re import split
 from typing import Any
+from typing import Callable
+from typing import cast
 from typing import Optional
 from typing import Sequence
 from typing import Tuple
+from typing import TypeVar
 from typing import Union
 
 from flask import current_app
@@ -30,6 +33,8 @@
 
 LocationType = Union[str, Sequence, None]
 
+_F = TypeVar("_F", bound=Callable[..., Any])
+
 
 def _verify_token_is_fresh(jwt_header: dict, jwt_data: dict) -> None:
     fresh = jwt_data["fresh"]
@@ -125,7 +130,7 @@ def jwt_required(
     locations: Optional[LocationType] = None,
     verify_type: bool = True,
     skip_revocation_check: bool = False,
-) -> Any:
+) -> Callable[[_F], _F]:
     """
     A decorator to protect a Flask endpoint with JSON Web Tokens.
 
@@ -161,15 +166,15 @@ def jwt_required(
         revocation status of the token will be checked.
     """
 
-    def wrapper(fn):
+    def wrapper(fn: _F) -> _F:
         @wraps(fn)
-        def decorator(*args, **kwargs):
+        def decorator(*args: Any, **kwargs: Any) -> Any:
             verify_jwt_in_request(
                 optional, fresh, refresh, locations, verify_type, skip_revocation_check
             )
             return current_app.ensure_sync(fn)(*args, **kwargs)
 
-        return decorator
+        return cast(_F, decorator)
 
     return wrapper
 

tests/test_typing.py

@@ -0,0 +1,138 @@
+"""
+Static typing checks for the public decorators.
+
+These are verified by mypy (see the ``mypy`` tox env), which fails if a
+decorator starts erasing the type of the function it wraps (e.g. by being
+annotated as returning ``Any``). The runtime tests make sure the decorated
+objects keep behaving like the originals.
+"""
+from typing import Any
+from typing import Callable
+from typing import Coroutine
+from typing import Dict
+from typing import Optional
+from typing import TYPE_CHECKING
+
+from flask import Flask
+from flask import Response
+
+from flask_jwt_extended import jwt_required
+from flask_jwt_extended import JWTManager
+
+if TYPE_CHECKING:
+    from typing_extensions import assert_type
+
+
+def test_jwt_required_preserves_function_types() -> None:
+    @jwt_required()
+    def view(user_id: str) -> Dict[str, str]:
+        return {"user_id": user_id}
+
+    @jwt_required(optional=True)
+    async def async_view() -> str:
+        return "ok"
+
+    if TYPE_CHECKING:
+        assert_type(view(""), Dict[str, str])
+        assert_type(async_view, Callable[[], Coroutine[Any, Any, str]])
+
+    assert view.__name__ == "view"
+    assert async_view.__name__ == "async_view"
+
+
+def test_loaders_preserve_function_types() -> None:
+    app = Flask(__name__)
+    jwt = JWTManager(app)
+
+    @jwt.additional_claims_loader
+    def additional_claims(identity: int) -> Dict[str, int]:
+        return {"identity": identity}
+
+    @jwt.additional_headers_loader
+    def additional_headers(identity: int) -> Dict[str, str]:
+        return {"kid": str(identity)}
+
+    @jwt.decode_key_loader
+    def decode_key(jwt_header: dict, jwt_data: dict) -> str:
+        return "decode-key"
+
+    @jwt.encode_key_loader
+    def encode_key(identity: int) -> str:
+        return "encode-key"
+
+    @jwt.expired_token_loader
+    def expired_token(jwt_header: dict, jwt_data: dict) -> Response:
+        return Response("expired", status=401)
+
+    @jwt.invalid_token_loader
+    def invalid_token(reason: str) -> Response:
+        return Response(reason, status=422)
+
+    @jwt.needs_fresh_token_loader
+    def needs_fresh_token(jwt_header: dict, jwt_data: dict) -> Response:
+        return Response("needs fresh", status=401)
+
+    @jwt.revoked_token_loader
+    def revoked_token(jwt_header: dict, jwt_data: dict) -> Response:
+        return Response("revoked", status=401)
+
+    @jwt.token_in_blocklist_loader
+    def in_blocklist(jwt_header: dict, jwt_data: dict) -> bool:
+        return False
+
+    @jwt.token_verification_failed_loader
+    def verification_failed(jwt_header: dict, jwt_data: dict) -> Response:
+        return Response("verification failed", status=400)
+
+    @jwt.token_verification_loader
+    def verify_token(jwt_header: dict, jwt_data: dict) -> bool:
+        return True
+
+    @jwt.unauthorized_loader
+    def unauthorized(reason: str) -> Response:
+        return Response(reason, status=401)
+
+    @jwt.user_identity_loader
+    def identity(user: int) -> str:
+        return str(user)
+
+    @jwt.user_lookup_loader
+    def user_lookup(jwt_header: dict, jwt_data: dict) -> Optional[int]:
+        return 1
+
+    @jwt.user_lookup_error_loader
+    def user_lookup_error(jwt_header: dict, jwt_data: dict) -> Response:
+        return Response("user not found", status=401)
+
+    if TYPE_CHECKING:
+        assert_type(additional_claims(1), Dict[str, int])
+        assert_type(additional_headers(1), Dict[str, str])
+        assert_type(decode_key({}, {}), str)
+        assert_type(encode_key(1), str)
+        assert_type(expired_token({}, {}), Response)
+        assert_type(invalid_token(""), Response)
+        assert_type(needs_fresh_token({}, {}), Response)
+        assert_type(revoked_token({}, {}), Response)
+        assert_type(in_blocklist({}, {}), bool)
+        assert_type(verification_failed({}, {}), Response)
+        assert_type(verify_token({}, {}), bool)
+        assert_type(unauthorized(""), Response)
+        assert_type(identity(1), str)
+        assert_type(user_lookup({}, {}), Optional[int])
+        assert_type(user_lookup_error({}, {}), Response)
+
+    assert additional_claims(1) == {"identity": 1}
+    assert additional_headers(1) == {"kid": "1"}
+    assert decode_key({}, {}) == "decode-key"
+    assert encode_key(1) == "encode-key"
+    assert expired_token({}, {}).status_code == 401
+    assert invalid_token("bad").status_code == 422
+    assert needs_fresh_token({}, {}).status_code == 401
+    assert revoked_token({}, {}).status_code == 401
+    assert in_blocklist({}, {}) is False
+    assert verification_failed({}, {}).status_code == 400
+    assert verify_token({}, {}) is True
+    assert unauthorized("missing").status_code == 401
+    assert identity(1) == "1"
+    assert user_lookup({}, {}) == 1
+    assert user_lookup_error({}, {}).status_code == 401