From c2be2dcdb709fe7a42006c90e95c1af6255c2301 Mon Sep 17 00:00:00 2001 From: Guido Iaquinti Date: Thu, 1 Oct 2026 17:53:22 +0200 Subject: [PATCH 1/2] Preserve types of functions decorated with jwt_required and JWTManager loaders --- flask_jwt_extended/jwt_manager.py | 33 ++++++++++-------- flask_jwt_extended/view_decorators.py | 13 ++++--- tests/test_typing.py | 49 +++++++++++++++++++++++++++ 3 files changed, 76 insertions(+), 19 deletions(-) create mode 100644 tests/test_typing.py diff --git a/flask_jwt_extended/jwt_manager.py b/flask_jwt_extended/jwt_manager.py index 7f0d967..e6002c9 100644 --- a/flask_jwt_extended/jwt_manager.py +++ b/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 diff --git a/flask_jwt_extended/view_decorators.py b/flask_jwt_extended/view_decorators.py index b7dbca0..5afc283 100644 --- a/flask_jwt_extended/view_decorators.py +++ b/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 diff --git a/tests/test_typing.py b/tests/test_typing.py new file mode 100644 index 0000000..950049f --- /dev/null +++ b/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 From 395602a516ddc0a76f378897a1937d9c436b4d0d Mon Sep 17 00:00:00 2001 From: Guido Iaquinti Date: Sun, 4 Oct 2026 16:41:05 +0200 Subject: [PATCH 2/2] Fix CI and address CR comments --- docs/conf.py | 2 + tests/test_typing.py | 115 ++++++++++++++++++++++++++++++++++++++----- 2 files changed, 104 insertions(+), 13 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index b5dd348..d503004 100644 --- a/docs/conf.py +++ b/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"), ] diff --git a/tests/test_typing.py b/tests/test_typing.py index 950049f..6197340 100644 --- a/tests/test_typing.py +++ b/tests/test_typing.py @@ -3,14 +3,18 @@ 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 +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 @@ -19,10 +23,7 @@ from typing_extensions import assert_type -def test_decorators_preserve_function_types() -> None: - app = Flask(__name__) - jwt = JWTManager(app) - +def test_jwt_required_preserves_function_types() -> None: @jwt_required() def view(user_id: str) -> Dict[str, str]: return {"user_id": user_id} @@ -31,19 +32,107 @@ def view(user_id: str) -> Dict[str, str]: async def async_view() -> str: return "ok" - @jwt.user_identity_loader - def identity(user: int) -> str: - return str(user) + 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(view(""), Dict[str, str]) - assert_type(identity(1), str) + 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 view.__name__ == "view" - assert identity(1) == "1" + 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