Experimental: FSM Middleware.

aiogram/contrib/middlewares/fsm.py

@@ -0,0 +1,80 @@
+import copy
+import weakref
+
+from aiogram.dispatcher.middlewares import LifetimeControllerMiddleware
+from aiogram.dispatcher.storage import FSMContext
+
+
+class FSMMiddleware(LifetimeControllerMiddleware):
+    skip_patterns = ['error', 'update']
+
+    def __init__(self):
+        super(FSMMiddleware, self).__init__()
+        self._proxies = weakref.WeakKeyDictionary()
+
+    async def pre_process(self, obj, data, *args):
+        proxy = await FSMSStorageProxy.create(self.manager.dispatcher.current_state())
+        data['state_data'] = proxy
+
+    async def post_process(self, obj, data, *args):
+        proxy = data.get('state_data', None)
+        if isinstance(proxy, FSMSStorageProxy):
+            await proxy.save()
+
+
+class FSMSStorageProxy(dict):
+    def __init__(self, fsm_context: FSMContext):
+        super(FSMSStorageProxy, self).__init__()
+        self.fsm_context = fsm_context
+        self._copy = {}
+        self._data = {}
+        self._state = None
+        self._is_dirty = False
+
+    @classmethod
+    async def create(cls, fsm_context: FSMContext):
+        """
+        :param fsm_context:
+        :return:
+        """
+        proxy = cls(fsm_context)
+        await proxy.load()
+        return proxy
+
+    async def load(self):
+        self.clear()
+        self._state = await self.fsm_context.get_state()
+        self.update(await self.fsm_context.get_data())
+        self._copy = copy.deepcopy(self)
+        self._is_dirty = False
+
+    @property
+    def state(self):
+        return self._state
+
+    @state.setter
+    def state(self, value):
+        self._state = value
+        self._is_dirty = True
+
+    @state.deleter
+    def state(self):
+        self._state = None
+        self._is_dirty = True
+
+    async def save(self, force=False):
+        if self._copy != self or force:
+            await self.fsm_context.set_data(data=self)
+        if self._is_dirty or force:
+            await self.fsm_context.set_state(self.state)
+        self._is_dirty = False
+        self._copy = copy.deepcopy(self)
+
+    def __str__(self):
+        s = super(FSMSStorageProxy, self).__str__()
+        readable_state = f"'{self.state}'" if self.state else "''"
+        return f"<{self.__class__.__name__}(state={readable_state}, data={s})>"
+
+    def clear(self):
+        del self.state
+        return super(FSMSStorageProxy, self).clear()

aiogram/dispatcher/middlewares.py

@@ -101,3 +101,28 @@ class BaseMiddleware:
         if not handler:
             return None
         await handler(*args)
+
+
+class LifetimeControllerMiddleware(BaseMiddleware):
+    # TODO: Rename class
+
+    skip_patterns = None
+
+    async def pre_process(self, obj, data, *args):
+        pass
+
+    async def post_process(self, obj, data, *args):
+        pass
+
+    async def trigger(self, action, args):
+        if self.skip_patterns is not None and any(item in action for item in self.skip_patterns):
+            return False
+
+        obj, *args, data = args
+        if action.startswith('pre_process_'):
+            await self.pre_process(obj, data, *args)
+        elif action.startswith('post_process_'):
+            await self.post_process(obj, data, *args)
+        else:
+            return False
+        return True