From d389b5f9cbab6fdef67a6efcf46d3d2c20c4f8bf Mon Sep 17 00:00:00 2001
From: Atsu Kakitani <atsunori.kakitani@gmail.com>
Date: Wed, 18 May 2022 23:19:49 -0700
Subject: [PATCH] Import SQLAlchemy types in Flask-SQLAlchemy (#7861)

---
 stubs/Flask-SQLAlchemy/METADATA.toml            |  1 +
 .../flask_sqlalchemy/__init__.pyi               | 17 +++++++++--------
 .../Flask-SQLAlchemy/flask_sqlalchemy/model.pyi |  4 ++--
 tests/pytype_exclude_list.txt                   |  3 +++
 4 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/stubs/Flask-SQLAlchemy/METADATA.toml b/stubs/Flask-SQLAlchemy/METADATA.toml
index 43ea0392c..1fafe940f 100644
--- a/stubs/Flask-SQLAlchemy/METADATA.toml
+++ b/stubs/Flask-SQLAlchemy/METADATA.toml
@@ -1 +1,2 @@
 version = "2.5.*"
+requires = ["types-SQLAlchemy"]
diff --git a/stubs/Flask-SQLAlchemy/flask_sqlalchemy/__init__.pyi b/stubs/Flask-SQLAlchemy/flask_sqlalchemy/__init__.pyi
index c8cbad071..7d7713ee6 100644
--- a/stubs/Flask-SQLAlchemy/flask_sqlalchemy/__init__.pyi
+++ b/stubs/Flask-SQLAlchemy/flask_sqlalchemy/__init__.pyi
@@ -1,19 +1,18 @@
-from typing import Any
+from typing import Any, TypeVar
+
+from sqlalchemy.orm.query import Query
+from sqlalchemy.orm.session import Session
 
 from . import utils as utils
 from .model import DefaultMeta as DefaultMeta, Model as Model
 
-# SQLAlchemy is not part of typeshed
-_Query = Any
-_SessionBase = Any
-
 models_committed: Any
 before_models_committed: Any
 
-class SignallingSession(_SessionBase):
+class SignallingSession(Session):
     app: Any
     def __init__(self, db, autocommit: bool = ..., autoflush: bool = ..., **options) -> None: ...
-    def get_bind(self, mapper: Any | None = ..., clause: Any | None = ...): ...
+    def get_bind(self, mapper: Any | None = ..., clause: Any | None = ...): ...  # type: ignore[override]
 
 def get_debug_queries(): ...
 
@@ -40,7 +39,9 @@ class Pagination:
         self, left_edge: int = ..., left_current: int = ..., right_current: int = ..., right_edge: int = ...
     ) -> None: ...
 
-class BaseQuery(_Query):
+_T = TypeVar("_T")
+
+class BaseQuery(Query[_T]):
     def get_or_404(self, ident, description: Any | None = ...): ...
     def first_or_404(self, description: Any | None = ...): ...
     def paginate(
diff --git a/stubs/Flask-SQLAlchemy/flask_sqlalchemy/model.pyi b/stubs/Flask-SQLAlchemy/flask_sqlalchemy/model.pyi
index 7909be376..24aec4589 100644
--- a/stubs/Flask-SQLAlchemy/flask_sqlalchemy/model.pyi
+++ b/stubs/Flask-SQLAlchemy/flask_sqlalchemy/model.pyi
@@ -1,6 +1,6 @@
 from typing import Any, Pattern
 
-_DeclarativeMeta = Any  # SQLAlchemy is not part of typeshed
+from sqlalchemy.ext.declarative import DeclarativeMeta
 
 def should_set_tablename(cls: type) -> bool: ...
 
@@ -15,7 +15,7 @@ class NameMetaMixin(type):
 class BindMetaMixin(type):
     def __init__(cls, name, bases, d) -> None: ...
 
-class DefaultMeta(NameMetaMixin, BindMetaMixin, _DeclarativeMeta): ...
+class DefaultMeta(NameMetaMixin, BindMetaMixin, DeclarativeMeta): ...
 
 class Model:
     query_class: Any | None
diff --git a/tests/pytype_exclude_list.txt b/tests/pytype_exclude_list.txt
index 75740bafa..4d91f5873 100644
--- a/tests/pytype_exclude_list.txt
+++ b/tests/pytype_exclude_list.txt
@@ -108,3 +108,6 @@ stubs/SQLAlchemy/sqlalchemy/sql/functions.pyi
 stubs/SQLAlchemy/sqlalchemy/sql/lambdas.pyi
 stubs/SQLAlchemy/sqlalchemy/sql/schema.pyi
 stubs/SQLAlchemy/sqlalchemy/sql/selectable.pyi
+stubs/Flask-SQLAlchemy/flask_sqlalchemy/__init__.pyi
+stubs/Flask-SQLAlchemy/flask_sqlalchemy/model.pyi
+stubs/Flask-SQLAlchemy/flask_sqlalchemy/utils.pyi