rename some ref for corrent type hint

Author: Zoupers

src/flask_sqlalchemy/cli.py

@@ -5,12 +5,18 @@
 from flask import current_app
 
 
+if t.TYPE_CHECKING:
+    from .extension import SQLAlchemy as SQLAlchemy
+else:
+    SQLAlchemy = t.Any
+
+
 def add_models_to_shell() -> dict[str, t.Any]:
     """Registered with :meth:`~flask.Flask.shell_context_processor` if
     ``add_models_to_shell`` is enabled. Adds the ``db`` instance and all model classes
     to ``flask shell``.
     """
-    db = current_app.extensions["sqlalchemy"]
+    db: SQLAlchemy = current_app.extensions["sqlalchemy"]
     out = {m.class_.__name__: m.class_ for m in db.Model._sa_registry.mappers}
     out["db"] = db
     return out

src/flask_sqlalchemy/extension.py

@@ -127,7 +127,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] | sqlalchemy.orm.DeclarativeMeta = Model,
         engine_options: dict[str, t.Any] | None = None,
         add_models_to_shell: bool = True,
     ):
@@ -338,7 +338,7 @@ def init_app(self, app: Flask) -> None:
 
     def _make_scoped_session(
         self, options: dict[str, t.Any]
-    ) -> sa.orm.scoped_session[Session]:
+    ) -> sqlalchemy.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`.
 
@@ -361,11 +361,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 sqlalchemy.orm.scoped_session(factory, scope)
 
     def _make_session_factory(
         self, options: dict[str, t.Any]
-    ) -> sa.orm.sessionmaker[Session]:
+    ) -> sqlalchemy.orm.sessionmaker[Session]:
         """Create the SQLAlchemy :class:`sqlalchemy.orm.sessionmaker` used by
         :meth:`_make_scoped_session`.
 
@@ -388,7 +388,7 @@ def _make_session_factory(
         """
         options.setdefault("class_", Session)
         options.setdefault("query_cls", self.Query)
-        return sa.orm.sessionmaker(db=self, **options)
+        return sqlalchemy.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
@@ -482,7 +482,7 @@ def __new__(
         return Table
 
     def _make_declarative_base(
-        self, model: type[Model] | sa.orm.DeclarativeMeta
+        self, model: type[Model] | sqlalchemy.orm.DeclarativeMeta
     ) -> type[t.Any]:
         """Create a SQLAlchemy declarative model class. The result is available as
         :attr:`Model`.
@@ -503,9 +503,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, sqlalchemy.orm.DeclarativeMeta):
             metadata = self._make_metadata(None)
-            model = sa.orm.declarative_base(
+            model = sqlalchemy.orm.declarative_base(
                 metadata=metadata, cls=model, name="Model", metaclass=DefaultMeta
             )
 
@@ -780,7 +780,7 @@ def one_or_404(
         """
         try:
             return self.session.execute(statement).scalar_one()
-        except (sa.exc.NoResultFound, sa.exc.MultipleResultsFound):
+        except (sqlalchemy.exc.NoResultFound, sqlalchemy.exc.MultipleResultsFound):
             abort(404, description=description)
 
     def paginate(
@@ -859,7 +859,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 sqlalchemy.exc.UnboundExecutionError(message) from None
 
             metadata = self.metadatas[key]
             getattr(metadata, op_name)(bind=engine)
@@ -936,31 +936,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]:
+    ) -> sqlalchemy.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 sqlalchemy.orm.relationship(*args, **kwargs)
 
     def dynamic_loader(
         self, argument: t.Any, **kwargs: t.Any
-    ) -> sa.orm.RelationshipProperty[t.Any]:
+    ) -> sqlalchemy.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 sqlalchemy.orm.dynamic_loader(argument, **kwargs)
 
     def _relation(
         self, *args: t.Any, **kwargs: t.Any
-    ) -> sa.orm.RelationshipProperty[t.Any]:
+    ) -> sqlalchemy.orm.RelationshipProperty[t.Any]:
         """A :func:`sqlalchemy.orm.relationship` that applies this extension's
         :attr:`Query` class for dynamic relationships and backrefs.
 
@@ -973,7 +973,7 @@ def _relation(
         """
         # Deprecated, removed in SQLAlchemy 2.0. Accessed through ``__getattr__``.
         self._set_rel_query(kwargs)
-        f = sa.orm.relation  # type: ignore[attr-defined]
+        f = sqlalchemy.orm.relation  # type: ignore[attr-defined]
         return f(*args, **kwargs)  # type: ignore[no-any-return]
 
     def __getattr__(self, name: str) -> t.Any:
@@ -993,12 +993,12 @@ def __getattr__(self, name: str) -> t.Any:
             return self._relation
 
         if name == "event":
-            return sa.event
+            return sqlalchemy.event
 
         if name.startswith("_"):
             raise AttributeError(name)
 
-        for mod in (sa, sa.orm):
+        for mod in (sa, sqlalchemy.orm):
             if hasattr(mod, name):
                 return getattr(mod, name)
 

src/flask_sqlalchemy/model.py

@@ -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, sqlalchemy.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__"], sqlalchemy.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, sqlalchemy.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, sqlalchemy.orm.DeclarativeMeta):
     """SQLAlchemy declarative metaclass that provides ``__bind_key__`` and
     ``__tablename__`` support.
     """

src/flask_sqlalchemy/pagination.py

@@ -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(sqlalchemy.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,7 +2,6 @@
 
 import typing as t
 
-import sqlalchemy as sa
 import sqlalchemy.exc
 import sqlalchemy.orm
 from flask import abort
@@ -11,7 +10,7 @@
 from .pagination import QueryPagination
 
 
-class Query(sa.orm.Query):  # type: ignore[type-arg]
+class Query(sqlalchemy.orm.Query[t.Any]):
     """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 (sqlalchemy.exc.NoResultFound, sqlalchemy.exc.MultipleResultsFound):
             abort(404, description=description)
 
     def paginate(

src/flask_sqlalchemy/record_queries.py

@@ -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)
+    sqlalchemy.event.listen(engine, "before_cursor_execute", _record_start, named=True)
+    sqlalchemy.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

@@ -11,7 +11,7 @@
     from .extension import SQLAlchemy
 
 
-class Session(sa.orm.Session):
+class Session(sqlalchemy.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 sqlalchemy.exc.NoInspectionAvailable as e:
                 if isinstance(mapper, type):
-                    raise sa.orm.exc.UnmappedClassError(mapper) from e
+                    raise sqlalchemy.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 sqlalchemy.exc.UnboundExecutionError(
                 f"Bind key '{key}' is not in 'SQLALCHEMY_BINDS' config."
             )
 

src/flask_sqlalchemy/track_modifications.py

@@ -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: sqlalchemy.orm.scoped_session[Session]) -> None:
+    sqlalchemy.event.listen(session, "before_flush", _record_ops, named=True)
+    sqlalchemy.event.listen(session, "before_commit", _record_ops, named=True)
+    sqlalchemy.event.listen(session, "before_commit", _before_commit)
+    sqlalchemy.event.listen(session, "after_commit", _after_commit)
+    sqlalchemy.event.listen(session, "after_rollback", _after_rollback)
 
 
 def _record_ops(session: Session, **kwargs: t.Any) -> None:

tests/test_legacy_query.py

@@ -15,9 +15,9 @@
 
 @pytest.fixture(autouse=True)
 def ignore_query_warning() -> t.Generator[None, None, None]:
-    if hasattr(sa.exc, "LegacyAPIWarning"):
+    if hasattr(sqlalchemy.exc, "LegacyAPIWarning"):
         with warnings.catch_warnings():
-            exc = sa.exc.LegacyAPIWarning
+            exc = sqlalchemy.exc.LegacyAPIWarning
             warnings.simplefilter("ignore", exc)
             yield
     else:

tests/test_metadata.py

@@ -26,15 +26,15 @@ def test_custom_metadata() -> None:
 
 
 def test_metadata_from_custom_model() -> None:
-    base = sa.orm.declarative_base(cls=Model, metaclass=DefaultMeta)
+    base = sqlalchemy.orm.declarative_base(cls=Model, metaclass=DefaultMeta)
     metadata = base.metadata
     db = SQLAlchemy(model_class=base)
     assert db.Model.metadata is metadata
     assert db.Model.metadata is db.metadata
 
 
 def test_custom_metadata_overrides_custom_model() -> None:
-    base = sa.orm.declarative_base(cls=Model, metaclass=DefaultMeta)
+    base = sqlalchemy.orm.declarative_base(cls=Model, metaclass=DefaultMeta)
     metadata = sa.MetaData()
     db = SQLAlchemy(model_class=base, metadata=metadata)
     assert db.Model.metadata is metadata
@@ -69,21 +69,21 @@ class Post(db.Model):
         __bind_key__ = "a"
         id = sa.Column(sa.Integer, primary_key=True)
 
-    with pytest.raises(sa.exc.OperationalError):
+    with pytest.raises(sqlalchemy.exc.OperationalError):
         db.session.execute(sa.select(User)).scalars()
 
-    with pytest.raises(sa.exc.OperationalError):
+    with pytest.raises(sqlalchemy.exc.OperationalError):
         db.session.execute(sa.select(Post)).scalars()
 
     db.create_all()
     db.session.execute(sa.select(User)).scalars()
     db.session.execute(sa.select(Post)).scalars()
     db.drop_all()
 
-    with pytest.raises(sa.exc.OperationalError):
+    with pytest.raises(sqlalchemy.exc.OperationalError):
         db.session.execute(sa.select(User)).scalars()
 
-    with pytest.raises(sa.exc.OperationalError):
+    with pytest.raises(sqlalchemy.exc.OperationalError):
         db.session.execute(sa.select(Post)).scalars()
 
 
@@ -103,7 +103,7 @@ class Post(db.Model):
     db.create_all(bind_key=bind_key)
     db.session.execute(sa.select(Post)).scalars()
 
-    with pytest.raises(sa.exc.OperationalError):
+    with pytest.raises(sqlalchemy.exc.OperationalError):
         db.session.execute(sa.select(User)).scalars()