type cheking

src/core/db_manager.py

@@ -1,20 +1,24 @@
+from typing import Any
+
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
 from src.repository.tasks import TasksRepo
 from src.repository.users import UsersRepo
 
 
 class DBManager:
-    def __init__(self, session_factory):
+    def __init__(self, session_factory: async_sessionmaker[AsyncSession]):
         self.session_factory = session_factory
 
-    async def __aenter__(self):
-        self.session = self.session_factory()
+    async def __aenter__(self) -> "DBManager":
+        self.session: AsyncSession = self.session_factory()
         self.user = UsersRepo(self.session)
         self.task = TasksRepo(self.session)
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         await self.session.rollback()
         await self.session.close()
 
-    async def commit(self):
+    async def commit(self) -> None:
         await self.session.commit()

src/core/interfaces.py

@@ -1,15 +1,18 @@
-from typing import Protocol
+from typing import Any, Protocol
+
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from src.repository.tasks import TasksRepo
 from src.repository.users import UsersRepo
 
 
 class IUOWDB(Protocol):
+    session: AsyncSession
     user: UsersRepo
     task: TasksRepo
 
     async def __aenter__(self) -> "IUOWDB": ...
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: ...
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: ...
 
     async def commit(self) -> None: ...