Merge branch '3.0.x' into #1205

Author: pamelafox

CHANGES.rst

@@ -1,3 +1,11 @@
+Version 3.0.4
+-------------
+
+Unreleased
+
+-   Fix type hint for ``get_or_404`` return value. :pr:`1208`
+
+
 Version 3.0.3
 -------------
 

src/flask_sqlalchemy/extension.py

@@ -23,6 +23,8 @@
 from .session import Session
 from .table import _Table
 
+_O = t.TypeVar("_O", bound=object)  # Based on sqlalchemy.orm._typing.py
+
 
 class SQLAlchemy:
     """Integrates SQLAlchemy with Flask. This handles setting up one or more engines,
@@ -731,8 +733,8 @@ def get_binds(self) -> dict[sa.Table, sa.engine.Engine]:
         }
 
     def get_or_404(
-        self, entity: type[t.Any], ident: t.Any, *, description: str | None = None
-    ) -> t.Any:
+        self, entity: type[_O], ident: t.Any, *, description: str | None = None
+    ) -> t.Optional[_O]:
         """Like :meth:`session.get() <sqlalchemy.orm.Session.get>` but aborts with a
         ``404 Not Found`` error instead of returning ``None``.
 

tests/conftest.py

@@ -31,7 +31,7 @@ def db(app: Flask) -> SQLAlchemy:
 
 
 @pytest.fixture
-def Todo(app: Flask, db: SQLAlchemy) -> t.Any:
+def Todo(app: Flask, db: SQLAlchemy) -> t.Generator[t.Any, None, None]:
     class Todo(db.Model):
         id = sa.Column(sa.Integer, primary_key=True)
         title = sa.Column(sa.String)

tests/test_view_query.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import typing as t
+
+import pytest
+import sqlalchemy as sa
+from flask import Flask
+from werkzeug.exceptions import NotFound
+
+from flask_sqlalchemy import SQLAlchemy
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_view_get_or_404(db: SQLAlchemy, Todo: t.Any) -> None:
+    item = Todo()
+    db.session.add(item)
+    db.session.commit()
+    assert db.get_or_404(Todo, 1) is item
+    with pytest.raises(NotFound):
+        assert db.get_or_404(Todo, 2)
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_first_or_404(db: SQLAlchemy, Todo: t.Any) -> None:
+    db.session.add(Todo(title="a"))
+    db.session.commit()
+    result = db.first_or_404(db.select(Todo).filter_by(title="a"))
+    assert result.title == "a"
+
+    with pytest.raises(NotFound):
+        db.first_or_404(db.select(Todo).filter_by(title="b"))
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_view_one_or_404(db: SQLAlchemy, Todo: t.Any) -> None:
+    db.session.add(Todo(title="a"))
+    db.session.add(Todo(title="b"))
+    db.session.add(Todo(title="b"))
+    db.session.commit()
+    result = db.one_or_404(db.select(Todo).filter_by(title="a"))
+    assert result.title == "a"
+
+    with pytest.raises(NotFound):
+        # MultipleResultsFound
+        db.one_or_404(db.select(Todo).filter_by(title="b"))
+
+    with pytest.raises(NotFound):
+        # NoResultFound
+        db.one_or_404(db.select(Todo).filter_by(title="c"))
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_paginate(db: SQLAlchemy, Todo: t.Any) -> None:
+    db.session.add_all(Todo() for _ in range(150))
+    db.session.commit()
+    p = db.paginate(db.select(Todo))
+    assert p.total == 150
+    assert len(p.items) == 20
+    p2 = p.next()
+    assert p2.page == 2
+    assert p2.total == 150
+
+
+# This test creates its own inline model so that it can use that as the type
+@pytest.mark.usefixtures("app_ctx")
+def test_view_get_or_404_typed(db: SQLAlchemy, app: Flask) -> None:
+    class Quiz(db.Model):
+        id = sa.Column(sa.Integer, primary_key=True)
+        topic = sa.Column(sa.String)
+
+    db.create_all()
+
+    item: Quiz = Quiz()
+    db.session.add(item)
+    db.session.commit()
+    result = db.get_or_404(Quiz, 1)
+    assert result is item
+    if hasattr(t, "assert_type"):
+        t.assert_type(result, Quiz)
+    with pytest.raises(NotFound):
+        assert db.get_or_404(Quiz, 2)
+    db.drop_all()

tox.ini

@@ -29,7 +29,9 @@ wheel_build_env = .pkg
 constrain_package_deps = true
 use_frozen_constraints = true
 deps = -r requirements/mypy.txt
-commands = mypy
+commands =
+    mypy --python-version 3.7
+    mypy --python-version 3.11
 
 [testenv:docs]
 package = wheel