Fix little problem with type hints and adopt a test to make sure it works in all platform (#1207)

Author: pamelafox

src/flask_sqlalchemy/extension.py

@@ -5,9 +5,9 @@
 from weakref import WeakKeyDictionary
 
 import sqlalchemy as sa
-import sqlalchemy.event
-import sqlalchemy.exc
-import sqlalchemy.orm
+import sqlalchemy.event as sa_event
+import sqlalchemy.exc as sa_exc
+import sqlalchemy.orm as sa_orm
 from flask import abort
 from flask import current_app
 from flask import Flask
@@ -129,7 +129,7 @@ def __init__(
         metadata: sa.MetaData | None = None,
         session_options: dict[str, t.Any] | None = None,
         query_class: type[Query] = Query,
-        model_class: type[Model] | sa.orm.DeclarativeMeta = Model,
+        model_class: type[Model] | sa_orm.DeclarativeMeta = Model,
         engine_options: dict[str, t.Any] | None = None,
         add_models_to_shell: bool = True,
     ):
@@ -341,7 +341,7 @@ def init_app(self, app: Flask) -> None:
 
     def _make_scoped_session(
         self, options: dict[str, t.Any]
-    ) -> sa.orm.scoped_session[Session]:
+    ) -> sa_orm.scoped_session[Session]:
         """Create a :class:`sqlalchemy.orm.scoping.scoped_session` around the factory
         from :meth:`_make_session_factory`. The result is available as :attr:`session`.
 
@@ -364,11 +364,11 @@ def _make_scoped_session(
         """
         scope = options.pop("scopefunc", _app_ctx_id)
         factory = self._make_session_factory(options)
-        return sa.orm.scoped_session(factory, scope)
+        return sa_orm.scoped_session(factory, scope)
 
     def _make_session_factory(
         self, options: dict[str, t.Any]
-    ) -> sa.orm.sessionmaker[Session]:
+    ) -> sa_orm.sessionmaker[Session]:
         """Create the SQLAlchemy :class:`sqlalchemy.orm.sessionmaker` used by
         :meth:`_make_scoped_session`.
 
@@ -391,7 +391,7 @@ def _make_session_factory(
         """
         options.setdefault("class_", Session)
         options.setdefault("query_cls", self.Query)
-        return sa.orm.sessionmaker(db=self, **options)
+        return sa_orm.sessionmaker(db=self, **options)
 
     def _teardown_commit(self, exc: BaseException | None) -> None:
         """Commit the session at the end of the request if there was not an unhandled
@@ -485,7 +485,7 @@ def __new__(
         return Table
 
     def _make_declarative_base(
-        self, model: type[Model] | sa.orm.DeclarativeMeta
+        self, model: type[Model] | sa_orm.DeclarativeMeta
     ) -> type[t.Any]:
         """Create a SQLAlchemy declarative model class. The result is available as
         :attr:`Model`.
@@ -506,9 +506,9 @@ def _make_declarative_base(
         .. versionchanged:: 2.3
             ``model`` can be an already created declarative model class.
         """
-        if not isinstance(model, sa.orm.DeclarativeMeta):
+        if not isinstance(model, sa_orm.DeclarativeMeta):
             metadata = self._make_metadata(None)
-            model = sa.orm.declarative_base(
+            model = sa_orm.declarative_base(
                 metadata=metadata, cls=model, name="Model", metaclass=DefaultMeta
             )
 
@@ -783,7 +783,7 @@ def one_or_404(
         """
         try:
             return self.session.execute(statement).scalar_one()
-        except (sa.exc.NoResultFound, sa.exc.MultipleResultsFound):
+        except (sa_exc.NoResultFound, sa_exc.MultipleResultsFound):
             abort(404, description=description)
 
     def paginate(
@@ -862,7 +862,7 @@ def _call_for_binds(
                 if key is None:
                     message = f"'SQLALCHEMY_DATABASE_URI' config is not set. {message}"
 
-                raise sa.exc.UnboundExecutionError(message) from None
+                raise sa_exc.UnboundExecutionError(message) from None
 
             metadata = self.metadatas[key]
             getattr(metadata, op_name)(bind=engine)
@@ -939,31 +939,31 @@ def _set_rel_query(self, kwargs: dict[str, t.Any]) -> None:
 
     def relationship(
         self, *args: t.Any, **kwargs: t.Any
-    ) -> sa.orm.RelationshipProperty[t.Any]:
+    ) -> sa_orm.RelationshipProperty[t.Any]:
         """A :func:`sqlalchemy.orm.relationship` that applies this extension's
         :attr:`Query` class for dynamic relationships and backrefs.
 
         .. versionchanged:: 3.0
             The :attr:`Query` class is set on ``backref``.
         """
         self._set_rel_query(kwargs)
-        return sa.orm.relationship(*args, **kwargs)
+        return sa_orm.relationship(*args, **kwargs)
 
     def dynamic_loader(
         self, argument: t.Any, **kwargs: t.Any
-    ) -> sa.orm.RelationshipProperty[t.Any]:
+    ) -> sa_orm.RelationshipProperty[t.Any]:
         """A :func:`sqlalchemy.orm.dynamic_loader` that applies this extension's
         :attr:`Query` class for relationships and backrefs.
 
         .. versionchanged:: 3.0
             The :attr:`Query` class is set on ``backref``.
         """
         self._set_rel_query(kwargs)
-        return sa.orm.dynamic_loader(argument, **kwargs)
+        return sa_orm.dynamic_loader(argument, **kwargs)
 
     def _relation(
         self, *args: t.Any, **kwargs: t.Any
-    ) -> sa.orm.RelationshipProperty[t.Any]:
+    ) -> sa_orm.RelationshipProperty[t.Any]:
         """A :func:`sqlalchemy.orm.relationship` that applies this extension's
         :attr:`Query` class for dynamic relationships and backrefs.
 
@@ -976,8 +976,8 @@ def _relation(
         """
         # Deprecated, removed in SQLAlchemy 2.0. Accessed through ``__getattr__``.
         self._set_rel_query(kwargs)
-        f = sa.orm.relation  # type: ignore[attr-defined]
-        return f(*args, **kwargs)  # type: ignore[no-any-return]
+        f = sa_orm.relationship
+        return f(*args, **kwargs)
 
     def __getattr__(self, name: str) -> t.Any:
         if name == "db":
@@ -996,12 +996,12 @@ def __getattr__(self, name: str) -> t.Any:
             return self._relation
 
         if name == "event":
-            return sa.event
+            return sa_event
 
         if name.startswith("_"):
             raise AttributeError(name)
 
-        for mod in (sa, sa.orm):
+        for mod in (sa, sa_orm):
             if hasattr(mod, name):
                 return getattr(mod, name)
 

src/flask_sqlalchemy/model.py

@@ -4,7 +4,7 @@
 import typing as t
 
 import sqlalchemy as sa
-import sqlalchemy.orm
+import sqlalchemy.orm as sa_orm
 
 from .query import Query
 
@@ -182,21 +182,21 @@ def should_set_tablename(cls: type) -> bool:
     joined-table inheritance. If no primary key is found, the name will be unset.
     """
     if cls.__dict__.get("__abstract__", False) or not any(
-        isinstance(b, sa.orm.DeclarativeMeta) for b in cls.__mro__[1:]
+        isinstance(b, sa_orm.DeclarativeMeta) for b in cls.__mro__[1:]
     ):
         return False
 
     for base in cls.__mro__:
         if "__tablename__" not in base.__dict__:
             continue
 
-        if isinstance(base.__dict__["__tablename__"], sa.orm.declared_attr):
+        if isinstance(base.__dict__["__tablename__"], sa_orm.declared_attr):
             return False
 
         return not (
             base is cls
             or base.__dict__.get("__abstract__", False)
-            or not isinstance(base, sa.orm.DeclarativeMeta)
+            or not isinstance(base, sa_orm.DeclarativeMeta)
         )
 
     return True
@@ -208,7 +208,7 @@ def camel_to_snake_case(name: str) -> str:
     return name.lower().lstrip("_")
 
 
-class DefaultMeta(BindMetaMixin, NameMetaMixin, sa.orm.DeclarativeMeta):
+class DefaultMeta(BindMetaMixin, NameMetaMixin, sa_orm.DeclarativeMeta):
     """SQLAlchemy declarative metaclass that provides ``__bind_key__`` and
     ``__tablename__`` support.
     """

src/flask_sqlalchemy/pagination.py

@@ -4,7 +4,7 @@
 from math import ceil
 
 import sqlalchemy as sa
-import sqlalchemy.orm
+import sqlalchemy.orm as sa_orm
 from flask import abort
 from flask import request
 
@@ -336,7 +336,7 @@ def _query_items(self) -> list[t.Any]:
 
     def _query_count(self) -> int:
         select = self._query_args["select"]
-        sub = select.options(sa.orm.lazyload("*")).order_by(None).subquery()
+        sub = select.options(sa_orm.lazyload("*")).order_by(None).subquery()
         session = self._query_args["session"]
         out = session.execute(sa.select(sa.func.count()).select_from(sub)).scalar()
         return out  # type: ignore[no-any-return]

src/flask_sqlalchemy/query.py

@@ -2,16 +2,15 @@
 
 import typing as t
 
-import sqlalchemy as sa
-import sqlalchemy.exc
-import sqlalchemy.orm
+import sqlalchemy.exc as sa_exc
+import sqlalchemy.orm as sa_orm
 from flask import abort
 
 from .pagination import Pagination
 from .pagination import QueryPagination
 
 
-class Query(sa.orm.Query):  # type: ignore[type-arg]
+class Query(sa_orm.Query):  # type: ignore[type-arg]
     """SQLAlchemy :class:`~sqlalchemy.orm.query.Query` subclass with some extra methods
     useful for querying in a web application.
 
@@ -58,7 +57,7 @@ def one_or_404(self, description: str | None = None) -> t.Any:
         """
         try:
             return self.one()
-        except (sa.exc.NoResultFound, sa.exc.MultipleResultsFound):
+        except (sa_exc.NoResultFound, sa_exc.MultipleResultsFound):
             abort(404, description=description)
 
     def paginate(

src/flask_sqlalchemy/record_queries.py

@@ -6,7 +6,7 @@
 from time import perf_counter
 
 import sqlalchemy as sa
-import sqlalchemy.event
+import sqlalchemy.event as sa_event
 from flask import current_app
 from flask import g
 from flask import has_app_context
@@ -96,8 +96,8 @@ def __getitem__(self, key: int) -> object:
 
 
 def _listen(engine: sa.engine.Engine) -> None:
-    sa.event.listen(engine, "before_cursor_execute", _record_start, named=True)
-    sa.event.listen(engine, "after_cursor_execute", _record_end, named=True)
+    sa_event.listen(engine, "before_cursor_execute", _record_start, named=True)
+    sa_event.listen(engine, "after_cursor_execute", _record_end, named=True)
 
 
 def _record_start(context: sa.engine.ExecutionContext, **kwargs: t.Any) -> None:

src/flask_sqlalchemy/session.py

@@ -3,15 +3,15 @@
 import typing as t
 
 import sqlalchemy as sa
-import sqlalchemy.exc
-import sqlalchemy.orm
+import sqlalchemy.exc as sa_exc
+import sqlalchemy.orm as sa_orm
 from flask.globals import app_ctx
 
 if t.TYPE_CHECKING:
     from .extension import SQLAlchemy
 
 
-class Session(sa.orm.Session):
+class Session(sa_orm.Session):
     """A SQLAlchemy :class:`~sqlalchemy.orm.Session` class that chooses what engine to
     use based on the bind key associated with the metadata associated with the thing
     being queried.
@@ -55,9 +55,9 @@ def get_bind(
         if mapper is not None:
             try:
                 mapper = sa.inspect(mapper)
-            except sa.exc.NoInspectionAvailable as e:
+            except sa_exc.NoInspectionAvailable as e:
                 if isinstance(mapper, type):
-                    raise sa.orm.exc.UnmappedClassError(mapper) from e
+                    raise sa_orm.exc.UnmappedClassError(mapper) from e
 
                 raise
 
@@ -88,7 +88,7 @@ def _clause_to_engine(
         key = clause.metadata.info["bind_key"]
 
         if key not in engines:
-            raise sa.exc.UnboundExecutionError(
+            raise sa_exc.UnboundExecutionError(
                 f"Bind key '{key}' is not in 'SQLALCHEMY_BINDS' config."
             )
 

src/flask_sqlalchemy/track_modifications.py

@@ -3,8 +3,8 @@
 import typing as t
 
 import sqlalchemy as sa
-import sqlalchemy.event
-import sqlalchemy.orm
+import sqlalchemy.event as sa_event
+import sqlalchemy.orm as sa_orm
 from flask import current_app
 from flask import has_app_context
 from flask.signals import Namespace  # type: ignore[attr-defined]
@@ -29,12 +29,12 @@
 """
 
 
-def _listen(session: sa.orm.scoped_session[Session]) -> None:
-    sa.event.listen(session, "before_flush", _record_ops, named=True)
-    sa.event.listen(session, "before_commit", _record_ops, named=True)
-    sa.event.listen(session, "before_commit", _before_commit)
-    sa.event.listen(session, "after_commit", _after_commit)
-    sa.event.listen(session, "after_rollback", _after_rollback)
+def _listen(session: sa_orm.scoped_session[Session]) -> None:
+    sa_event.listen(session, "before_flush", _record_ops, named=True)
+    sa_event.listen(session, "before_commit", _record_ops, named=True)
+    sa_event.listen(session, "before_commit", _before_commit)
+    sa_event.listen(session, "after_commit", _after_commit)
+    sa_event.listen(session, "after_rollback", _after_rollback)
 
 
 def _record_ops(session: Session, **kwargs: t.Any) -> None:

tests/test_legacy_query.py

@@ -5,7 +5,7 @@
 
 import pytest
 import sqlalchemy as sa
-import sqlalchemy.exc
+import sqlalchemy.exc as sa_exc
 from flask import Flask
 from werkzeug.exceptions import NotFound
 
@@ -15,9 +15,9 @@
 
 @pytest.fixture(autouse=True)
 def ignore_query_warning() -> t.Generator[None, None, None]:
-    if hasattr(sa.exc, "LegacyAPIWarning"):
+    if hasattr(sa_exc, "LegacyAPIWarning"):
         with warnings.catch_warnings():
-            exc = sa.exc.LegacyAPIWarning
+            exc = sa_exc.LegacyAPIWarning
             warnings.simplefilter("ignore", exc)
             yield
     else: