Define on_after_forgot_password with a decorator

docs/configuration/router.md

@@ -2,33 +2,13 @@
 
 We're almost there! The last step is to configure the `FastAPIUsers` object that will wire the database adapter, the authentication class and the user model to expose the FastAPI router.
 
-## Hooks
-
-In order to be as unopinionated as possible, you'll have to define your logic after some actions.
-
-### After forgot password
-
-This hook is called after a successful forgot password request. It is called with **two arguments**: the **user** which has requested to reset their password and a ready-to-use **JWT token** that will be accepted by the reset password route.
-
-Typically, you'll want to **send an e-mail** with the link (and the token) that allows the user to reset their password.
-
-You can define it as an `async` or standard method.
-
-Example:
-
-```py
-def on_after_forgot_password(user, token):
-    print(f'User {user.id} has forgot their password. Reset token: {token}')
-```
-
 ## Configure `FastAPIUsers`
 
-The last step is to instantiate `FastAPIUsers` object with all the elements we defined before. More precisely:
+Configure `FastAPIUsers` object with all the elements we defined before. More precisely:
 
 * `db`: Database adapter instance.
 * `auth`: Authentication logic instance.
 * `user_model`: Pydantic model of a user.
-* `on_after_forgot_password`: Hook called after a forgot password request.
 * `reset_password_token_secret`: Secret to encode reset password token.
 * `reset_password_token_lifetime_seconds`: Lifetime of reset password token in seconds. Default to one hour.
 
@@ -39,7 +19,6 @@ fastapi_users = FastAPIUsers(
     user_db,
     auth,
     User,
-    on_after_forgot_password,
     SECRET,
 )
 ```
@@ -51,6 +30,26 @@ app = FastAPI()
 app.include_router(fastapi_users.router, prefix="/users", tags=["users"])
 ```
 
+## Event handlers
+
+In order to be as unopinionated as possible, we expose decorators that allow you to plug your own logic after some actions. You can have several handlers per event.
+
+### After forgot password
+
+This event handler is called after a successful forgot password request. It is called with **two arguments**: the **user** which has requested to reset their password and a ready-to-use **JWT token** that will be accepted by the reset password route.
+
+Typically, you'll want to **send an e-mail** with the link (and the token) that allows the user to reset their password.
+
+You can define it as an `async` or standard method.
+
+Example:
+
+```py
+@fastapi_users.on_after_forgot_password()
+def on_after_forgot_password(user, token):
+    print(f'User {user.id} has forgot their password. Reset token: {token}')
+```
+
 ## Next steps
 
 Check out a [full example](full_example.md) that will show you the big picture.

docs/src/full_sqlalchemy.py

@@ -35,16 +35,16 @@ class User(BaseUser):
 
 auth = JWTAuthentication(secret=SECRET, lifetime_seconds=3600)
 
+app = FastAPI()
+fastapi_users = FastAPIUsers(user_db, auth, User, SECRET)
+app.include_router(fastapi_users.router, prefix="/users", tags=["users"])
 
+
+@fastapi_users.on_after_forgot_password()
 def on_after_forgot_password(user, token):
     print(f"User {user.id} has forgot their password. Reset token: {token}")
 
 
-app = FastAPI()
-fastapi_users = FastAPIUsers(user_db, auth, User, on_after_forgot_password, SECRET)
-app.include_router(fastapi_users.router, prefix="/users", tags=["users"])
-
-
 @app.on_event("startup")
 async def startup():
     await database.connect()

fastapi_users/fastapi_users.py

@@ -1,11 +1,9 @@
-from typing import Any, Callable, Type
+from typing import Callable, Type
-
-from fastapi import APIRouter
 
 from fastapi_users.authentication import BaseAuthentication
 from fastapi_users.db import BaseUserDatabase
 from fastapi_users.models import BaseUser, BaseUserDB
-from fastapi_users.router import get_user_router
+from fastapi_users.router import Events, UserRouter, get_user_router
 
 
 class FastAPIUsers:
@@ -15,17 +13,16 @@ class FastAPIUsers:
     :param db: Database adapter instance.
     :param auth: Authentication logic instance.
     :param user_model: Pydantic model of a user.
-    :param on_after_forgot_password: Hook called after a forgot password request.
     :param reset_password_token_secret: Secret to encode reset password token.
     :param reset_password_token_lifetime_seconds: Lifetime of reset password token.
 
-    :attribute router: FastAPI router exposing authentication routes.
+    :attribute router: Router exposing authentication routes.
     :attribute get_current_user: Dependency callable to inject authenticated user.
     """
 
     db: BaseUserDatabase
     auth: BaseAuthentication
-    router: APIRouter
+    router: UserRouter
     get_current_user: Callable[..., BaseUserDB]
 
     def __init__(
@@ -33,7 +30,6 @@ class FastAPIUsers:
         db: BaseUserDatabase,
         auth: BaseAuthentication,
         user_model: Type[BaseUser],
-        on_after_forgot_password: Callable[[BaseUserDB, str], Any],
         reset_password_token_secret: str,
         reset_password_token_lifetime_seconds: int = 3600,
     ):
@@ -43,7 +39,6 @@ class FastAPIUsers:
             self.db,
             user_model,
             self.auth,
-            on_after_forgot_password,
             reset_password_token_secret,
             reset_password_token_lifetime_seconds,
         )
@@ -56,3 +51,14 @@ class FastAPIUsers:
 
         get_current_superuser = self.auth.get_current_superuser(self.db)
         self.get_current_superuser = get_current_superuser  # type: ignore
+
+    def on_after_forgot_password(self) -> Callable:
+        """Add an event handler on successful forgot password request."""
+        return self._on_event(Events.ON_AFTER_FORGOT_PASSWORD)
+
+    def _on_event(self, event_type: Events) -> Callable:
+        def decorator(func: Callable) -> Callable:
+            self.router.add_event_handler(event_type, func)
+            return func
+
+        return decorator

fastapi_users/router.py

@@ -1,5 +1,7 @@
 import asyncio
-from typing import Any, Callable, Type
+import typing
+from collections import defaultdict
+from enum import Enum
 
 import jwt
 from fastapi import APIRouter, Body, Depends, HTTPException
@@ -10,27 +12,45 @@ from starlette.responses import Response
 
 from fastapi_users.authentication import BaseAuthentication
 from fastapi_users.db import BaseUserDatabase
-from fastapi_users.models import BaseUser, BaseUserDB, Models
+from fastapi_users.models import BaseUser, Models
 from fastapi_users.password import get_password_hash
 from fastapi_users.utils import JWT_ALGORITHM, generate_jwt
 
 
+class Events(Enum):
+    ON_AFTER_FORGOT_PASSWORD = 1
+
+
+class UserRouter(APIRouter):
+    event_handlers: typing.DefaultDict[Events, typing.List[typing.Callable]]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.event_handlers = defaultdict(list)
+
+    def add_event_handler(self, event_type: Events, func: typing.Callable) -> None:
+        self.event_handlers[event_type].append(func)
+
+    async def run_handlers(self, event_type: Events, *args, **kwargs) -> None:
+        for handler in self.event_handlers[event_type]:
+            if asyncio.iscoroutinefunction(handler):
+                await handler(*args, **kwargs)
+            else:
+                handler(*args, **kwargs)
+
+
 def get_user_router(
     user_db: BaseUserDatabase,
-    user_model: Type[BaseUser],
+    user_model: typing.Type[BaseUser],
     auth: BaseAuthentication,
-    on_after_forgot_password: Callable[[BaseUserDB, str], Any],
     reset_password_token_secret: str,
     reset_password_token_lifetime_seconds: int = 3600,
-) -> APIRouter:
+) -> UserRouter:
     """Generate a router with the authentication routes."""
-    router = APIRouter()
+    router = UserRouter()
     models = Models(user_model)
 
     reset_password_token_audience = "fastapi-users:reset"
-    is_on_after_forgot_password_async = asyncio.iscoroutinefunction(
-        on_after_forgot_password
-    )
 
     get_current_active_user = auth.get_current_active_user(user_db)
 
@@ -74,10 +94,7 @@ def get_user_router(
                 reset_password_token_lifetime_seconds,
                 reset_password_token_secret,
             )
-            if is_on_after_forgot_password_async:
+            await router.run_handlers(Events.ON_AFTER_FORGOT_PASSWORD, user, token)
-                await on_after_forgot_password(user, token)
-            else:
-                on_after_forgot_password(user, token)
 
         return None
 

tests/test_fastapi_users.py

@@ -22,9 +22,12 @@ def test_app_client(request, mock_user_db, mock_authentication) -> TestClient:
     class User(BaseUser):
         pass
 
-    fastapi_users = FastAPIUsers(
+    fastapi_users = FastAPIUsers(mock_user_db, mock_authentication, User, SECRET)
-        mock_user_db, mock_authentication, User, request.param, SECRET
+
-    )
+    @fastapi_users.on_after_forgot_password()
+    def on_after_forgot_password():
+        return request.param()
+
     app = FastAPI()
     app.include_router(fastapi_users.router)
 

tests/test_router.py

@@ -8,7 +8,7 @@ from starlette import status
 from starlette.testclient import TestClient
 
 from fastapi_users.models import BaseUser, BaseUserDB
-from fastapi_users.router import get_user_router
+from fastapi_users.router import Events, get_user_router
 from fastapi_users.utils import JWT_ALGORITHM, generate_jwt
 
 SECRET = "SECRET"
@@ -47,12 +47,11 @@ def test_app_client(
         pass
 
     userRouter = get_user_router(
-        mock_user_db,
+        mock_user_db, User, mock_authentication, SECRET, LIFETIME
-        User,
+    )
-        mock_authentication,
+
-        on_after_forgot_password,
+    userRouter.add_event_handler(
-        SECRET,
+        Events.ON_AFTER_FORGOT_PASSWORD, on_after_forgot_password
-        LIFETIME,
     )
 
     app = FastAPI()