Define on_after_forgot_password with a decorator

fastapi_users/fastapi_users.py

@@ -1,11 +1,9 @@
-from typing import Any, Callable, Type
-
-from fastapi import APIRouter
+from typing import Callable, Type
 
 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