diff --git a/aiosqlite/cursor.py b/aiosqlite/cursor.py index 6932874..a2990ba 100644 --- a/aiosqlite/cursor.py +++ b/aiosqlite/cursor.py @@ -19,7 +19,7 @@ def __aiter__(self) -> AsyncIterator[sqlite3.Row]: """The cursor proxy is also an async iterator.""" return self._fetch_chunked() - async def _fetch_chunked(self): + async def _fetch_chunked(self) -> AsyncIterator[sqlite3.Row]: while True: rows = await self.fetchmany(self.iter_chunk_size) if not rows: @@ -56,14 +56,14 @@ async def fetchone(self) -> Optional[sqlite3.Row]: """Fetch a single row.""" return await self._execute(self._cursor.fetchone) - async def fetchmany(self, size: Optional[int] = None) -> Iterable[sqlite3.Row]: + async def fetchmany(self, size: Optional[int] = None) -> list[sqlite3.Row]: """Fetch up to `cursor.arraysize` number of rows.""" args: tuple[int, ...] = () if size is not None: args = (size,) return await self._execute(self._cursor.fetchmany, *args) - async def fetchall(self) -> Iterable[sqlite3.Row]: + async def fetchall(self) -> list[sqlite3.Row]: """Fetch all remaining rows.""" return await self._execute(self._cursor.fetchall) @@ -103,8 +103,8 @@ def row_factory(self, factory: Optional[type]) -> None: def connection(self) -> sqlite3.Connection: return self._cursor.connection - async def __aenter__(self): + async def __aenter__(self) -> "Cursor": return self - async def __aexit__(self, exc_type, exc_val, exc_tb): + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: await self.close() diff --git a/aiosqlite/tests/smoke.py b/aiosqlite/tests/smoke.py index 046f9b7..e8e5017 100644 --- a/aiosqlite/tests/smoke.py +++ b/aiosqlite/tests/smoke.py @@ -126,6 +126,37 @@ async def test_iterable_cursor(self): assert len(rows) == 10 + async def test_iterable_cursor_after_fetchmany(self): + async with aiosqlite.connect(self.db) as db: + cursor = await db.cursor() + await cursor.execute( + "create table iterable_cursor_after_fetchmany " + "(i integer primary key asc, k integer)" + ) + await cursor.executemany( + "insert into iterable_cursor_after_fetchmany (k) values (?)", + [[i] for i in range(5)], + ) + await db.commit() + + async with aiosqlite.connect(self.db) as db: + cursor = await db.execute( + "select * from iterable_cursor_after_fetchmany" + ) + first_two = await cursor.fetchmany(2) + self.assertEqual(len(first_two), 2) + self.assertEqual(first_two[0], (1, 0)) + self.assertEqual(first_two[1], (2, 1)) + + remaining = [] + async for row in cursor: + remaining.append(row) + + self.assertEqual(len(remaining), 3) + self.assertEqual(remaining[0], (3, 2)) + self.assertEqual(remaining[1], (4, 3)) + self.assertEqual(remaining[2], (5, 4)) + async def test_multi_loop_usage(self): results = {}