Skip to content

Commit

Permalink
add sqlalchemy namespace to flask shell
Browse files Browse the repository at this point in the history
  • Loading branch information
yuxiaoy1 committed Jun 24, 2024
1 parent 14144b5 commit 4f316c4
Show file tree
Hide file tree
Showing 3 changed files with 5 additions and 2 deletions.
2 changes: 1 addition & 1 deletion CHANGES.rst
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ Version 3.1.2
-------------

- Fix issue with calling ``repr()`` on ``SQLAlchemy`` instance with no default engine. :issue:`1295`

- Adds ``sqlalchemy``` namespace as ``sa``` to ``flask shell```. :pr:`1350`

Version 3.1.1
-------------
Expand Down
4 changes: 3 additions & 1 deletion src/flask_sqlalchemy/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,17 @@

import typing as t

import sqlalchemy as sa
from flask import current_app


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``.
to ``flask shell``. Adds the ``sqlalchemy`` namespace as ``sa`` to ``flask shell``.
"""
db = current_app.extensions["sqlalchemy"]
out = {m.class_.__name__: m.class_ for m in db.Model._sa_registry.mappers}
out["sa"] = sa
out["db"] = db
return out
1 change: 1 addition & 0 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,5 +11,6 @@
@pytest.mark.usefixtures("app_ctx")
def test_shell_context(db: SQLAlchemy, Todo: t.Any) -> None:
context = add_models_to_shell()
assert "sa" in context
assert context["db"] is db
assert context["Todo"] is Todo

0 comments on commit 4f316c4

Please sign in to comment.