Add type annotations to CursorWrapper fetch methods

Explicit fetchone, fetchmany, and fetchall methods with return type
annotations so type checkers know the return types. Also fixes several
places that assumed fetchone() always returns a row.

Author: davegaeddert

plain-models/plain/models/backends/base/features.py

@@ -184,6 +184,8 @@ def supports_transactions(self) -> bool:
             self.connection.rollback()
             self.connection.set_autocommit(True)
             cursor.execute("SELECT COUNT(X) FROM ROLLBACK_TEST")
-            (count,) = cursor.fetchone()
+            row = cursor.fetchone()
+            assert row is not None
+            (count,) = row
             cursor.execute("DROP TABLE ROLLBACK_TEST")
         return count == 0

plain-models/plain/models/backends/mysql/base.py

@@ -357,6 +357,7 @@ def mysql_server_data(self) -> dict[str, Any]:
             """
             )
             row = cursor.fetchone()
+        assert row is not None
         return {
             "version": row[0],
             "sql_mode": row[1],

plain-models/plain/models/backends/sqlite3/base.py

@@ -230,7 +230,9 @@ def disable_constraint_checking(self) -> bool:
             # Foreign key constraints cannot be turned off while in a multi-
             # statement transaction. Fetch the current state of the pragma
             # to determine if constraints are effectively disabled.
-            enabled = cursor.execute("PRAGMA foreign_keys").fetchone()[0]
+            row = cursor.execute("PRAGMA foreign_keys").fetchone()
+            assert row is not None
+            enabled = row[0]
         return not bool(enabled)
 
     def enable_constraint_checking(self) -> None:
@@ -273,10 +275,12 @@ def check_constraints(self, table_names: list[str] | None = None) -> None:
                 assert primary_key_column_name is not None, (
                     f"Table {table_name} must have a primary key"
                 )
-                primary_key_value, bad_value = cursor.execute(
+                row = cursor.execute(
                     f"SELECT {self.ops.quote_name(primary_key_column_name)}, {self.ops.quote_name(column_name)} FROM {self.ops.quote_name(table_name)} WHERE rowid = %s",
                     (rowid,),
                 ).fetchone()
+                assert row is not None
+                primary_key_value, bad_value = row
                 raise IntegrityError(
                     f"The row in table '{table_name}' with primary key '{primary_key_value}' has an "
                     f"invalid foreign key: {table_name}.{column_name} contains a value '{bad_value}' that "

plain-models/plain/models/backends/sqlite3/introspection.py

@@ -338,14 +338,11 @@ def get_constraints(
         """
         constraints: dict[str, dict[str, Any]] = {}
         # Find inline check constraints.
-        try:
-            table_schema = cursor.execute(
-                f"SELECT sql FROM sqlite_master WHERE type='table' and name={self.connection.ops.quote_name(table_name)}"
-            ).fetchone()[0]
-        except TypeError:
-            # table_name is a view.
-            pass
-        else:
+        row = cursor.execute(
+            f"SELECT sql FROM sqlite_master WHERE type='table' and name={self.connection.ops.quote_name(table_name)}"
+        ).fetchone()
+        if row is not None:
+            table_schema = row[0]
             columns = {
                 info.name for info in self.get_table_description(cursor, table_name)
             }

plain-models/plain/models/backends/sqlite3/schema.py

@@ -168,9 +168,9 @@ def alter_field(
                 # that don't affect the on-disk content.
                 # https://sqlite.org/lang_altertable.html#otheralter
                 with self.connection.cursor() as cursor:
-                    schema_version = cursor.execute("PRAGMA schema_version").fetchone()[
-                        0
-                    ]
+                    row = cursor.execute("PRAGMA schema_version").fetchone()
+                    assert row is not None
+                    schema_version = row[0]
                     cursor.execute("PRAGMA writable_schema = 1")
                     references_template = f' REFERENCES "{table_name}" ("%s") '
                     new_column_name = new_field.get_attname_column()[1]

plain-models/plain/models/backends/utils.py

@@ -61,7 +61,7 @@ def fetchone(self) -> tuple[Any, ...] | None:
         """Fetch the next row of a query result set."""
         ...
 
-    def fetchmany(self, size: int = 0) -> list[tuple[Any, ...]]:
+    def fetchmany(self, size: int | None = None) -> list[tuple[Any, ...]]:
         """Fetch the next set of rows of a query result set."""
         ...
 
@@ -79,7 +79,7 @@ def __init__(self, cursor: DBAPICursor, db: BaseDatabaseWrapper) -> None:
         self.cursor = cursor
         self.db = db
 
-    WRAP_ERROR_ATTRS = frozenset(["fetchone", "fetchmany", "fetchall", "nextset"])
+    WRAP_ERROR_ATTRS = frozenset(["nextset"])
 
     def __getattr__(self, attr: str) -> Any:
         cursor_attr = getattr(self.cursor, attr)
@@ -92,6 +92,20 @@ def __iter__(self) -> Iterator[tuple[Any, ...]]:
         with self.db.wrap_database_errors:
             yield from self.cursor
 
+    def fetchone(self) -> tuple[Any, ...] | None:
+        with self.db.wrap_database_errors:
+            return self.cursor.fetchone()
+
+    def fetchmany(self, size: int | None = None) -> list[tuple[Any, ...]]:
+        with self.db.wrap_database_errors:
+            if size is None:
+                return self.cursor.fetchmany()
+            return self.cursor.fetchmany(size)
+
+    def fetchall(self) -> list[tuple[Any, ...]]:
+        with self.db.wrap_database_errors:
+            return self.cursor.fetchall()
+
     def __enter__(self) -> Self:
         return self