Import SQLAlchemy types in Flask-SQLAlchemy (#7861)

Author: akakitani

stubs/Flask-SQLAlchemy/METADATA.toml

@@ -1 +1,2 @@
 version = "2.5.*"
+requires = ["types-SQLAlchemy"]

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(

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

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