Skip to content

Commit c5e5fe5

Browse files
ya7010ya7010
committed
test: add cursor type tests.
1 parent fb452db commit c5e5fe5

16 files changed

+395
-120
lines changed

‎turu-snowflake/src/turu/snowflake/async_connection.py‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,12 @@
66
import turu.core.cursor
77
import turu.core.mock
88
from turu.core.cursor import GenericNewRowType
9-
from turu.snowflake.cursor import GenericNewPandasDataFlame, GenericNewPyArrowTable
9+
from turu.snowflake.cursor import GenericNewPandasDataFrame, GenericNewPyArrowTable
1010
from turu.snowflake.features import (
1111
GenericNewPanderaDataFrameModel,
12+
PandasDataFrame,
1213
PanderaDataFrame,
14+
PyArrowTable,
1315
)
1416
from typing_extensions import Never, Self, Unpack, override
1517

@@ -102,7 +104,7 @@ async def rollback(self) -> None:
102104
self._raw_connection.rollback()
103105

104106
@override
105-
async def cursor(self) -> AsyncCursor[Never, Never, Never]:
107+
async def cursor(self) -> AsyncCursor[Never, PandasDataFrame, PyArrowTable]:
106108
return AsyncCursor(self._raw_connection.cursor())
107109

108110
@override
@@ -112,7 +114,7 @@ async def execute(
112114
parameters: Optional[Any] = None,
113115
/,
114116
**options: Unpack[ExecuteOptions],
115-
) -> AsyncCursor[Tuple[Any], Never, Never]:
117+
) -> AsyncCursor[Tuple[Any], PandasDataFrame, PyArrowTable]:
116118
"""Prepare and execute a database operation (query or command).
117119
118120
This is not defined in [PEP 249](https://peps.python.org/pep-0249/),
@@ -136,7 +138,7 @@ async def executemany(
136138
seq_of_parameters: Sequence[Any],
137139
/,
138140
**options: Unpack[ExecuteOptions],
139-
) -> AsyncCursor[Tuple[Any], Never, Never]:
141+
) -> AsyncCursor[Tuple[Any], PandasDataFrame, PyArrowTable]:
140142
"""Prepare a database operation (query or command)
141143
and then execute it against all parameter sequences or mappings.
142144
@@ -170,12 +172,12 @@ async def execute_map(
170172
@overload
171173
async def execute_map(
172174
self,
173-
row_type: Type[GenericNewPandasDataFlame],
175+
row_type: Type[GenericNewPandasDataFrame],
174176
operation: str,
175177
parameters: "Optional[Any]" = None,
176178
/,
177179
**options: Unpack[ExecuteOptions],
178-
) -> "AsyncCursor[Never, GenericNewPandasDataFlame, Never]":
180+
) -> "AsyncCursor[Never, GenericNewPandasDataFrame, Never]":
179181
...
180182

181183
@overload
@@ -206,7 +208,7 @@ async def execute_map(
206208
row_type: Union[
207209
Type[GenericNewRowType],
208210
Type[GenericNewPanderaDataFrameModel],
209-
Type[GenericNewPandasDataFlame],
211+
Type[GenericNewPandasDataFrame],
210212
Type[GenericNewPyArrowTable],
211213
],
212214
operation: str,
@@ -265,12 +267,12 @@ async def executemany_map(
265267
@overload
266268
async def executemany_map(
267269
self,
268-
row_type: Type[GenericNewPandasDataFlame],
270+
row_type: Type[GenericNewPandasDataFrame],
269271
operation: str,
270272
seq_of_parameters: Sequence[Any],
271273
/,
272274
**options: Unpack[ExecuteOptions],
273-
) -> AsyncCursor[Never, GenericNewPandasDataFlame, Never]:
275+
) -> AsyncCursor[Never, GenericNewPandasDataFrame, Never]:
274276
...
275277

276278
@overload
@@ -290,7 +292,7 @@ async def executemany_map(
290292
row_type: Union[
291293
Type[GenericNewRowType],
292294
Type[GenericNewPanderaDataFrameModel],
293-
Type[GenericNewPandasDataFlame],
295+
Type[GenericNewPandasDataFrame],
294296
Type[GenericNewPyArrowTable],
295297
],
296298
operation: str,

‎turu-snowflake/src/turu/snowflake/async_cursor.py‎

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -20,17 +20,19 @@
2020
import turu.snowflake.record.async_record_cursor
2121
from turu.core.cursor import map_row
2222
from turu.snowflake.cursor import (
23-
GenericNewPandasDataFlame,
23+
GenericNewPandasDataFrame,
2424
GenericNewPyArrowTable,
2525
GenericNewRowType,
26-
GenericPandasDataFlame,
26+
GenericPandasDataFrame,
2727
GenericPyArrowTable,
2828
GenericRowType,
2929
)
3030
from turu.snowflake.features import (
3131
GenericNewPanderaDataFrameModel,
32+
PandasDataFrame,
3233
PanderaDataFrame,
3334
PanderaDataFrameModel,
35+
PyArrowTable,
3436
)
3537
from typing_extensions import Never, Self, Unpack, override
3638

@@ -45,7 +47,7 @@ class ExecuteOptions(TypedDict, total=False):
4547

4648

4749
class AsyncCursor(
48-
Generic[GenericRowType, GenericPandasDataFlame, GenericPyArrowTable],
50+
Generic[GenericRowType, GenericPandasDataFrame, GenericPyArrowTable],
4951
turu.core.async_cursor.AsyncCursor[GenericRowType, Any],
5052
):
5153
def __init__(
@@ -80,7 +82,7 @@ async def execute(
8082
parameters: Optional[Any] = None,
8183
/,
8284
**options: Unpack[ExecuteOptions],
83-
) -> "AsyncCursor[Tuple[Any], Never, Never]":
85+
) -> "AsyncCursor[Tuple[Any], PandasDataFrame, PyArrowTable]":
8486
"""Prepare and execute a database operation (query or command).
8587
8688
Parameters:
@@ -104,7 +106,7 @@ async def executemany(
104106
seq_of_parameters: Sequence[Any],
105107
/,
106108
**options: Unpack[ExecuteOptions],
107-
) -> "AsyncCursor[Tuple[Any], Never, Never]":
109+
) -> "AsyncCursor[Tuple[Any], PandasDataFrame, PyArrowTable]":
108110
"""Prepare a database operation (query or command)
109111
and then execute it against all parameter sequences or mappings.
110112
@@ -150,12 +152,12 @@ async def execute_map(
150152
@overload
151153
async def execute_map(
152154
self,
153-
row_type: Type[GenericNewPandasDataFlame],
155+
row_type: Type[GenericNewPandasDataFrame],
154156
operation: str,
155157
parameters: "Optional[Any]" = None,
156158
/,
157159
**options: Unpack[ExecuteOptions],
158-
) -> "AsyncCursor[Never, GenericNewPandasDataFlame, Never]":
160+
) -> "AsyncCursor[Never, GenericNewPandasDataFrame, Never]":
159161
...
160162

161163
@overload
@@ -175,7 +177,7 @@ async def execute_map(
175177
row_type: Union[
176178
Type[GenericNewRowType],
177179
Type[GenericNewPanderaDataFrameModel],
178-
Type[GenericNewPandasDataFlame],
180+
Type[GenericNewPandasDataFrame],
179181
Type[GenericNewPyArrowTable],
180182
],
181183
operation: str,
@@ -226,12 +228,12 @@ async def executemany_map(
226228
@overload
227229
async def executemany_map(
228230
self,
229-
row_type: Type[GenericNewPandasDataFlame],
231+
row_type: Type[GenericNewPandasDataFrame],
230232
operation: str,
231233
seq_of_parameters: "Sequence[Any]",
232234
/,
233235
**options: Unpack[ExecuteOptions],
234-
) -> "AsyncCursor[Never, GenericNewPandasDataFlame, Never]":
236+
) -> "AsyncCursor[Never, GenericNewPandasDataFrame, Never]":
235237
...
236238

237239
@overload
@@ -251,7 +253,7 @@ async def executemany_map(
251253
row_type: Union[
252254
Type[GenericNewRowType],
253255
Type[GenericNewPanderaDataFrameModel],
254-
Type[GenericNewPandasDataFlame],
256+
Type[GenericNewPandasDataFrame],
255257
Type[GenericNewPyArrowTable],
256258
],
257259
operation: str,
@@ -318,21 +320,21 @@ async def fetch_arrow_batches(self) -> Iterator[GenericPyArrowTable]:
318320
Iterator[GenericPyArrowTable], self._raw_cursor.fetch_arrow_batches()
319321
)
320322

321-
async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFlame:
323+
async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
322324
"""Fetch Pandas dataframes."""
323325

324326
df = self._raw_cursor.fetch_pandas_all(**kwargs)
325327

326328
if self._row_type and issubclass(self._row_type, PanderaDataFrameModel):
327329
df = self._row_type.validate(df) # type: ignore[union-attr]
328330

329-
return cast(GenericPandasDataFlame, df)
331+
return cast(GenericPandasDataFrame, df)
330332

331-
async def fetch_pandas_batches(self, **kwargs) -> Iterator[GenericPandasDataFlame]:
333+
async def fetch_pandas_batches(self, **kwargs) -> Iterator[GenericPandasDataFrame]:
332334
"""Fetch Pandas dataframes in batches, where 'batch' refers to Snowflake Chunk."""
333335

334336
return cast(
335-
Iterator[GenericPandasDataFlame],
337+
Iterator[GenericPandasDataFrame],
336338
self._raw_cursor.fetch_pandas_batches(**kwargs),
337339
)
338340

‎turu-snowflake/src/turu/snowflake/connection.py‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@
99
from turu.core.cursor import GenericNewRowType
1010
from turu.snowflake.features import (
1111
GenericNewPanderaDataFrameModel,
12+
PandasDataFrame,
1213
PanderaDataFrame,
14+
PyArrowTable,
1315
)
1416
from typing_extensions import Never, Unpack, override
1517

@@ -18,7 +20,7 @@
1820
from .cursor import (
1921
Cursor,
2022
ExecuteOptions,
21-
GenericNewPandasDataFlame,
23+
GenericNewPandasDataFrame,
2224
GenericNewPyArrowTable,
2325
Self,
2426
)
@@ -108,7 +110,7 @@ def rollback(self) -> None:
108110
self._raw_connection.rollback()
109111

110112
@override
111-
def cursor(self) -> Cursor[Never, Never, Never]:
113+
def cursor(self) -> Cursor[Never, PandasDataFrame, PyArrowTable]:
112114
return Cursor(self._raw_connection.cursor())
113115

114116
@override
@@ -118,7 +120,7 @@ def execute(
118120
parameters: Optional[Any] = None,
119121
/,
120122
**options: Unpack[ExecuteOptions],
121-
) -> Cursor[Tuple[Any], Never, Never]:
123+
) -> Cursor[Tuple[Any], PandasDataFrame, PyArrowTable]:
122124
"""Prepare and execute a database operation (query or command).
123125
124126
This is not defined in [PEP 249](https://peps.python.org/pep-0249/),
@@ -142,7 +144,7 @@ def executemany(
142144
seq_of_parameters: Sequence[Any],
143145
/,
144146
**options: Unpack[ExecuteOptions],
145-
) -> Cursor[Tuple[Any], Never, Never]:
147+
) -> Cursor[Tuple[Any], PandasDataFrame, PyArrowTable]:
146148
"""Prepare a database operation (query or command)
147149
and then execute it against all parameter sequences or mappings.
148150
@@ -185,12 +187,12 @@ def execute_map(
185187
@overload
186188
def execute_map(
187189
self,
188-
row_type: Type[GenericNewPandasDataFlame],
190+
row_type: Type[GenericNewPandasDataFrame],
189191
operation: str,
190192
parameters: "Optional[Any]" = None,
191193
/,
192194
**options: Unpack[ExecuteOptions],
193-
) -> Cursor[Never, GenericNewPandasDataFlame, Never]:
195+
) -> Cursor[Never, GenericNewPandasDataFrame, Never]:
194196
...
195197

196198
@overload
@@ -210,7 +212,7 @@ def execute_map(
210212
row_type: Union[
211213
Type[GenericNewRowType],
212214
Type[GenericNewPanderaDataFrameModel],
213-
Type[GenericNewPandasDataFlame],
215+
Type[GenericNewPandasDataFrame],
214216
Type[GenericNewPyArrowTable],
215217
],
216218
operation: str,
@@ -269,12 +271,12 @@ def executemany_map(
269271
@overload
270272
def executemany_map(
271273
self,
272-
row_type: Type[GenericNewPandasDataFlame],
274+
row_type: Type[GenericNewPandasDataFrame],
273275
operation: str,
274276
seq_of_parameters: Sequence[Any],
275277
/,
276278
**options: Unpack[ExecuteOptions],
277-
) -> Cursor[Never, GenericNewPandasDataFlame, Never]:
279+
) -> Cursor[Never, GenericNewPandasDataFrame, Never]:
278280
...
279281

280282
@overload
@@ -294,7 +296,7 @@ def executemany_map(
294296
row_type: Union[
295297
Type[GenericNewRowType],
296298
Type[GenericNewPanderaDataFrameModel],
297-
Type[GenericNewPandasDataFlame],
299+
Type[GenericNewPandasDataFrame],
298300
Type[GenericNewPyArrowTable],
299301
],
300302
operation: str,

0 commit comments

Comments
 (0)