Allow overriding how the locale is set

fastapi_babel/middleware.py

@@ -1,51 +1,37 @@
 import re
+from pathlib import Path
+from typing import Optional, Callable
 from fastapi import Request, Response
 from fastapi.templating import Jinja2Templates
-from starlette.middleware.base import BaseHTTPMiddleware
-from starlette.middleware.base import RequestResponseEndpoint
-from starlette.middleware.base import DispatchFunction
 from starlette.types import ASGIApp
-from typing import Optional, Callable
-from .core import Babel
-from .local_context import context_var
-from .properties import RootConfigs
-from pathlib import Path
+from starlette.middleware.base import BaseHTTPMiddleware, DispatchFunction, RequestResponseEndpoint
 
+from .core import Babel, _context_var
+from .properties import RootConfigs
 
 LANGUAGES_PATTERN = re.compile(r"([a-z]{2})-?([A-Z]{2})?(;q=\d.\d{1,3})?")
 
-
 class BabelMiddleware(BaseHTTPMiddleware):
     def __init__(
         self,
         app: ASGIApp,
         babel_configs: RootConfigs,
-        locale_selector: Optional[Callable[[Request], Optional[str]]] = None,
         jinja2_templates: Optional[Jinja2Templates] = None,
         dispatch: Optional[DispatchFunction] = None,
+        locale_selector: Optional[Callable[[Request], Optional[str]]] = None,
     ) -> None:
+        """
+        :param locale_selector: a callable that takes the request and returns
+                                a locale string (e.g. 'de' or 'en'), or None if
+                                no override is desired.
+        """
         super().__init__(app, dispatch)
         self.babel_configs = babel_configs
         self.jinja2_templates = jinja2_templates
-        self.locale_selector = locale_selector or self._default_locale_selector
-
-    def _default_locale_selector(self, request: Request):
-        return request.headers.get("Accept-Language", None)
+        self.locale_selector = locale_selector
 
-    def get_language(self, babel: Babel, lang_code: Optional[str] = None):
-        """Applies an available language.
-
-        To apply an available language it will be searched in the language folder for an available one
-        and will also priotize the one with the highest quality value. The Fallback language will be the
-        taken from the BABEL_DEFAULT_LOCALE var.
-
-            Args:
-                babel (Babel): Request scoped Babel instance
-                lang_code (str): The Value of the Accept-Language Header.
-
-            Returns:
-                str: The language that should be used.
-        """
+    def get_language(self, babel: Babel, lang_code: str) -> str:
+        """Original Babel logic that parses Accept-Language."""
         if not lang_code:
             return babel.config.BABEL_DEFAULT_LOCALE
 
@@ -56,52 +42,46 @@ def get_language(self, babel: Babel, lang_code: Optional[str] = None):
         ]
         languages = sorted(
             languages, key=lambda x: x[1], reverse=True
-        )  # sort the priority, no priority comes last
+        )  # priority sort
         translation_directory = Path(babel.config.BABEL_TRANSLATION_DIRECTORY)
         translation_files = [i.name for i in translation_directory.iterdir()]
         explicit_priority = None
 
         for lang, quality in languages:
             if lang in translation_files:
-                if (
-                    not quality
-                ):  # languages without quality value having the highest priority 1
+                # no quality => highest priority
+                if not quality:
                     return lang
-
-                elif (
-                    not explicit_priority
-                ):  # set language with explicit priority <= priority 1
+                # remember the first we see with explicit priority
+                elif not explicit_priority:
                     explicit_priority = lang
 
-        # Return language with explicit priority or default value
-        return (
-            explicit_priority
-            if explicit_priority
-            else self.babel_configs.BABEL_DEFAULT_LOCALE
-        )
+        # fallback: either the explicit_priority or default locale
+        return explicit_priority or self.babel_configs.BABEL_DEFAULT_LOCALE
 
-    async def dispatch(
-        self, request: Request, call_next: RequestResponseEndpoint
-    ) -> Response:
-        """dispatch function
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        lang_code: Optional[str] = request.headers.get("Accept-Language", None)
 
-        Args:
-            request (Request): ...
-            call_next (RequestResponseEndpoint): ...
+        # Create a new Babel instance per request
+        request.state.babel = Babel(configs=self.babel_configs)
 
-        Returns:
-            Response: ...
-        """
-        lang_code: Optional[str] = self.locale_selector(request)
+        # 1) If a custom locale_selector is provided, call it first
+        override_locale = None
+        if self.locale_selector:
+            override_locale = self.locale_selector(request)
 
-        # Create a new Babel instance per request
-        babel = Babel(configs=self.babel_configs)
-        babel.locale = self.get_language(babel, lang_code)
-        context_var.set(babel.gettext)
-        if self.jinja2_templates:
-            babel.install_jinja(self.jinja2_templates)
+        # 2) If no override, fallback to the original logic
+        if override_locale:
+            request.state.babel.locale = override_locale
+        else:
+            request.state.babel.locale = self.get_language(request.state.babel, lang_code)
 
-        request.state.babel = babel
+        # set the context var for template usage
+        _context_var.set(request.state.babel.gettext)
+
+        # optionally install jinja support
+        if self.jinja2_templates:
+            request.state.babel.install_jinja(self.jinja2_templates)
 
         response: Response = await call_next(request)
         return response