Make transaction behaviour consistent between Sync and Async implementation

Author: febus982

sqlalchemy_bind_manager/_unit_of_work.py

@@ -1,6 +1,9 @@
 from contextlib import contextmanager, asynccontextmanager
 from typing import Iterable, Type, Iterator, AsyncIterator
 
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import Session
+
 from sqlalchemy_bind_manager import SQLAlchemyRepository, SQLAlchemyAsyncRepository
 from sqlalchemy_bind_manager._bind_manager import SQLAlchemyBind, SQLAlchemyAsyncBind
 from sqlalchemy_bind_manager._transaction_handler import (
@@ -21,9 +24,9 @@ def __init__(
             setattr(self, r.__name__, r(session=self._transaction_handler.session))
 
     @contextmanager
-    def transaction(self, read_only: bool = False) -> Iterator[None]:
-        with self._transaction_handler.get_session(read_only=read_only) as s:
-            yield s
+    def transaction(self, read_only: bool = False) -> Iterator[Session]:
+        with self._transaction_handler.get_session(read_only=read_only) as _s:
+            yield _s
 
 
 class AsyncUnitOfWork:
@@ -40,6 +43,6 @@ def __init__(
             setattr(self, r.__name__, r(session=self._transaction_handler.session))
 
     @asynccontextmanager
-    async def transaction(self, read_only: bool = False) -> AsyncIterator[None]:
-        async with self._transaction_handler.get_session(read_only=read_only):
-            yield None
+    async def transaction(self, read_only: bool = False) -> AsyncIterator[AsyncSession]:
+        async with self._transaction_handler.get_session(read_only=read_only) as _s:
+            yield _s