fix pyright warnings for BaseBatch. __aenter__

piccolo-orm · Jun 14, 2024 · 34e1261 · 34e1261
1 parent 2a2745e
commit 34e1261
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/piccolo/engine/base.py b/piccolo/engine/base.py
@@ -36,7 +36,7 @@ def validate_savepoint_name(savepoint_name: str) -> None:
 
 class BaseBatch(metaclass=ABCMeta):
     @abstractmethod
-    async def __aenter__(self, *args, **kwargs): ...
+    async def __aenter__(self: Self, *args, **kwargs) -> Self: ...
 
     @abstractmethod
     async def __aexit__(self, *args, **kwargs): ...

diff --git a/piccolo/engine/postgres.py b/piccolo/engine/postgres.py
@@ -64,7 +64,7 @@ async def __anext__(self) -> t.List[t.Dict]:
             raise StopAsyncIteration()
         return response
 
-    async def __aenter__(self):
+    async def __aenter__(self: Self) -> Self:
         self._transaction = self.connection.transaction()
         await self._transaction.start()
         querystring = self.query.querystrings[0]

diff --git a/piccolo/engine/sqlite.py b/piccolo/engine/sqlite.py
@@ -344,7 +344,7 @@ async def __anext__(self) -> t.List[t.Dict]:
             raise StopAsyncIteration()
         return response
 
-    async def __aenter__(self):
+    async def __aenter__(self: Self) -> Self:
         querystring = self.query.querystrings[0]
         template, template_args = querystring.compile_string()