1
0

middleware.py 3.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. """Middleware Mixin that allow to add middleware to the app."""
  2. from __future__ import annotations
  3. import asyncio
  4. import dataclasses
  5. from typing import List
  6. from reflex.event import Event
  7. from reflex.middleware import HydrateMiddleware, Middleware
  8. from reflex.state import BaseState, StateUpdate
  9. from .mixin import AppMixin
  10. @dataclasses.dataclass
  11. class MiddlewareMixin(AppMixin):
  12. """Middleware Mixin that allow to add middleware to the app."""
  13. # Middleware to add to the app. Users should use `add_middleware`. PRIVATE.
  14. middleware: List[Middleware] = dataclasses.field(default_factory=list)
  15. def _init_mixin(self):
  16. self.middleware.append(HydrateMiddleware())
  17. def add_middleware(self, middleware: Middleware, index: int | None = None):
  18. """Add middleware to the app.
  19. Args:
  20. middleware: The middleware to add.
  21. index: The index to add the middleware at.
  22. """
  23. if index is None:
  24. self.middleware.append(middleware)
  25. else:
  26. self.middleware.insert(index, middleware)
  27. async def _preprocess(self, state: BaseState, event: Event) -> StateUpdate | None:
  28. """Preprocess the event.
  29. This is where middleware can modify the event before it is processed.
  30. Each middleware is called in the order it was added to the app.
  31. If a middleware returns an update, the event is not processed and the
  32. update is returned.
  33. Args:
  34. state: The state to preprocess.
  35. event: The event to preprocess.
  36. Returns:
  37. An optional state to return.
  38. """
  39. for middleware in self.middleware:
  40. if asyncio.iscoroutinefunction(middleware.preprocess):
  41. out = await middleware.preprocess(app=self, state=state, event=event) # pyright: ignore [reportArgumentType]
  42. else:
  43. out = middleware.preprocess(app=self, state=state, event=event) # pyright: ignore [reportArgumentType]
  44. if out is not None:
  45. return out # pyright: ignore [reportReturnType]
  46. async def _postprocess(
  47. self, state: BaseState, event: Event, update: StateUpdate
  48. ) -> StateUpdate:
  49. """Postprocess the event.
  50. This is where middleware can modify the delta after it is processed.
  51. Each middleware is called in the order it was added to the app.
  52. Args:
  53. state: The state to postprocess.
  54. event: The event to postprocess.
  55. update: The current state update.
  56. Returns:
  57. The state update to return.
  58. """
  59. for middleware in self.middleware:
  60. if asyncio.iscoroutinefunction(middleware.postprocess):
  61. out = await middleware.postprocess(
  62. app=self, # pyright: ignore [reportArgumentType]
  63. state=state,
  64. event=event,
  65. update=update,
  66. )
  67. else:
  68. out = middleware.postprocess(
  69. app=self, # pyright: ignore [reportArgumentType]
  70. state=state,
  71. event=event,
  72. update=update,
  73. )
  74. if out is not None:
  75. return out # pyright: ignore [reportReturnType]
  76. return update