Правки по mypy

maxapi/context/__init__.py

@@ -1,6 +1,6 @@
 import asyncio
 
-from typing import Any, Dict
+from typing import Any, Dict, Optional, Union
 
 from ..context.state_machine import State, StatesGroup
 
@@ -19,7 +19,7 @@ class MemoryContext:
         self.chat_id = chat_id
         self.user_id = user_id
         self._context: Dict[str, Any] = {}
-        self._state: State | None = None
+        self._state: State | str | None = None
         self._lock = asyncio.Lock()
 
     async def get_data(self) -> dict[str, Any]:
@@ -58,7 +58,7 @@ class MemoryContext:
         async with self._lock:
             self._context.update(kwargs)
 
-    async def set_state(self, state: State | str = None):
+    async def set_state(self, state: Optional[Union[State, str]] = None):
         
         """
         Устанавливает новое состояние.