Add support of initializing the extension on a Blueprint

flask_inertia/inertia.py

@@ -34,7 +34,9 @@
 from http import HTTPStatus
 from typing import Any, Optional
 
-from flask import Flask, Response, current_app, request
+from flask import Blueprint, Flask, Response, current_app, request
+from flask.app import App
+from flask.blueprints import BlueprintSetupState
 from jinja2 import Template
 from jsmin import jsmin
 from markupsafe import Markup
@@ -47,12 +49,12 @@
 class Inertia:
     """Inertia Plugin for Flask."""
 
-    def __init__(self, app: Optional[Flask] = None):
+    def __init__(self, app: Optional[Flask | Blueprint] = None):
         self.app = None
         if app is not None:
             self.init_app(app)
 
-    def init_app(self, app: Flask):
+    def init_app(self, app: Flask | Blueprint):
         """Init as an app extension
 
         * Register before_request hook
@@ -61,13 +63,25 @@ def init_app(self, app: Flask):
         """
         self.app = app
         self._shared_data = {}
-        if not hasattr(app, "extensions"):
-            app.extensions = {}
-        app.extensions["inertia"] = self
+        if isinstance(app, Flask):
+            self._init_extension(app)
+        elif isinstance(app, Blueprint):
+            blueprint = app
+            # Register the extension once the blueprint is registered
+            blueprint.record_once(self.register_blueprint)
         app.context_processor(self.context_processor)
         app.before_request(self.process_incoming_inertia_requests)
         app.after_request(self.update_redirect)
 
+    def register_blueprint(self, state: BlueprintSetupState):
+        self._init_extension(state.app)
+
+    def _init_extension(self, app: App):
+        """Store a reference to the extension in the app's extensions."""
+        if not hasattr(app, "extensions"):
+            app.extensions = {}
+        app.extensions["inertia"] = self
+
     def process_incoming_inertia_requests(self) -> Optional[Response]:
         """Process incoming Inertia requests.
 

flask_inertia/version.py

@@ -33,14 +33,19 @@
 import hashlib
 import os
 
-from flask import current_app
+from flask import current_app, request
 
 
 def get_asset_version() -> str:
     """Calculate asset version to allow Inertia to automatically make a full page visit in case of changes."""
+    template_folder = current_app.template_folder
+    if request.blueprint is not None:
+        blueprint = current_app.blueprints[request.blueprint]
+        if blueprint.template_folder is not None:
+            template_folder = blueprint.template_folder
     template_path = os.path.join(
         current_app.root_path,
-        current_app.template_folder,
+        template_folder,
         current_app.config["INERTIA_TEMPLATE"],
     )
     with open(template_path, "rb") as template_file:

tests/python/test_app.py

@@ -28,7 +28,7 @@
 from http import HTTPStatus
 from unittest.mock import patch
 
-from flask import Flask, redirect, url_for
+from flask import Blueprint, Flask, redirect, url_for
 from parameterized import parameterized
 
 from flask_inertia import (
@@ -441,5 +441,32 @@ def test_shorthand_route(self):
         self.assertEqual(data.component, "FAQ")
 
 
+class TestBlueprint(unittest.TestCase):
+    """Flask-Inertia blueprint tests."""
+
+    def setUp(self):
+        self.blueprint = Blueprint("inertia", __name__, template_folder=".")
+        self.blueprint.add_url_rule("/", "index", index)
+
+        self.inertia = Inertia(self.blueprint)
+        self.inertia.add_shorthand_route("/faq/", "FAQ")
+
+        self.app = Flask(__name__, template_folder=".")
+        self.app.config.from_object(TestConfig)
+        self.app.register_blueprint(self.blueprint)
+
+        self.app.response_class = InertiaTestResponse
+        self.client = self.app.test_client()
+
+    def test_extension_registered(self):
+        self.assertIn("inertia", self.app.extensions)
+
+    def test_shorthand_route(self):
+        response = self.client.get("/faq/")
+        data = response.inertia("app")
+        self.assertEqual(response.status_code, HTTPStatus.OK)
+        self.assertEqual(data.component, "FAQ")
+
+
 if __name__ == "__main__":
     unittest.main()