Preserve types of functions decorated with jwt_required and loaders

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,49 @@
+"""
+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 test makes sure the decorated
+objects keep behaving like the originals.
+"""
+
+from typing import Dict
+from typing import TYPE_CHECKING
+
+from flask import Flask
+
+from flask_jwt_extended import jwt_required
+from flask_jwt_extended import JWTManager
+
+if TYPE_CHECKING:
+    from typing_extensions import assert_type
+
+
+def test_decorators_preserve_function_types() -> None:
+    app = Flask(__name__)
+    jwt = JWTManager(app)
+
+    @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"
+
+    @jwt.user_identity_loader
+    def identity(user: int) -> str:
+        return str(user)
+
+    @jwt.token_in_blocklist_loader
+    def in_blocklist(jwt_header: dict, jwt_data: dict) -> bool:
+        return False
+
+    if TYPE_CHECKING:
+        assert_type(view(""), Dict[str, str])
+        assert_type(identity(1), str)
+        assert_type(in_blocklist({}, {}), bool)
+
+    assert view.__name__ == "view"
+    assert identity(1) == "1"
+    assert in_blocklist({}, {}) is False