From 860db6f6c4a8c57d09de3be8d8d9805f403e985c Mon Sep 17 00:00:00 2001
From: Joseph Young <80432516+jpy-git@users.noreply.github.com>
Date: Mon, 14 Mar 2022 15:15:32 +0000
Subject: [PATCH] sqlalchemy: Annotate `(Async)Session.__(a)enter__` (#7488)

---
 stubs/SQLAlchemy/sqlalchemy/ext/asyncio/session.pyi | 3 ++-
 stubs/SQLAlchemy/sqlalchemy/orm/session.pyi         | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/stubs/SQLAlchemy/sqlalchemy/ext/asyncio/session.pyi b/stubs/SQLAlchemy/sqlalchemy/ext/asyncio/session.pyi
index 94b3d1bde..3b8a3c028 100644
--- a/stubs/SQLAlchemy/sqlalchemy/ext/asyncio/session.pyi
+++ b/stubs/SQLAlchemy/sqlalchemy/ext/asyncio/session.pyi
@@ -1,3 +1,4 @@
+from _typeshed import Self
 from typing import Any
 
 from ...util import memoized_property
@@ -50,7 +51,7 @@ class AsyncSession(ReversibleProxy):
     async def close(self): ...
     @classmethod
     async def close_all(cls): ...
-    async def __aenter__(self): ...
+    async def __aenter__(self: Self) -> Self: ...
     async def __aexit__(self, type_, value, traceback) -> None: ...
     # proxied from Session
     identity_map: Any
diff --git a/stubs/SQLAlchemy/sqlalchemy/orm/session.pyi b/stubs/SQLAlchemy/sqlalchemy/orm/session.pyi
index 68586fdf9..576b7b89b 100644
--- a/stubs/SQLAlchemy/sqlalchemy/orm/session.pyi
+++ b/stubs/SQLAlchemy/sqlalchemy/orm/session.pyi
@@ -1,3 +1,4 @@
+from _typeshed import Self
 from collections.abc import Mapping
 from typing import Any, TypeVar, overload
 
@@ -104,7 +105,7 @@ class Session(_SessionClassMethods):
         query_cls: Any | None = ...,
     ) -> None: ...
     connection_callable: Any
-    def __enter__(self): ...
+    def __enter__(self: Self) -> Self: ...
     def __exit__(self, type_, value, traceback) -> None: ...
     @property
     def transaction(self): ...