Skip to content

Commit c7c997a

Browse files
ya7010ya7010
authored
Merge pull request #62 from yassun7010/add_type_testcases
fix: execute() type hint.
2 parents 499598b + 4992944 commit c7c997a

16 files changed

+470
-195
lines changed

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

Lines changed: 21 additions & 19 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

@@ -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,44 +172,44 @@ 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
182184
async def execute_map(
183185
self,
184-
row_type: Type[GenericNewPanderaDataFrameModel],
186+
row_type: Type[GenericNewPyArrowTable],
185187
operation: str,
186188
parameters: "Optional[Any]" = None,
187189
/,
188190
**options: Unpack[ExecuteOptions],
189-
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
191+
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
190192
...
191193

192194
@overload
193195
async def execute_map(
194196
self,
195-
row_type: Type[GenericNewPyArrowTable],
197+
row_type: Type[GenericNewPanderaDataFrameModel],
196198
operation: str,
197199
parameters: "Optional[Any]" = None,
198200
/,
199201
**options: Unpack[ExecuteOptions],
200-
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
202+
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
201203
...
202204

203205
@override
204206
async def execute_map(
205207
self,
206208
row_type: Union[
207209
Type[GenericNewRowType],
208-
Type[GenericNewPanderaDataFrameModel],
209-
Type[GenericNewPandasDataFlame],
210+
Type[GenericNewPandasDataFrame],
210211
Type[GenericNewPyArrowTable],
212+
Type[GenericNewPanderaDataFrameModel],
211213
],
212214
operation: str,
213215
parameters: Optional[Any] = None,
@@ -254,44 +256,44 @@ async def executemany_map(
254256
@overload
255257
async def executemany_map(
256258
self,
257-
row_type: Type[GenericNewPanderaDataFrameModel],
259+
row_type: Type[GenericNewPandasDataFrame],
258260
operation: str,
259261
seq_of_parameters: Sequence[Any],
260262
/,
261263
**options: Unpack[ExecuteOptions],
262-
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
264+
) -> AsyncCursor[Never, GenericNewPandasDataFrame, Never]:
263265
...
264266

265267
@overload
266268
async def executemany_map(
267269
self,
268-
row_type: Type[GenericNewPandasDataFlame],
270+
row_type: Type[GenericNewPyArrowTable],
269271
operation: str,
270272
seq_of_parameters: Sequence[Any],
271273
/,
272274
**options: Unpack[ExecuteOptions],
273-
) -> AsyncCursor[Never, GenericNewPandasDataFlame, Never]:
275+
) -> AsyncCursor[Never, Never, GenericNewPyArrowTable]:
274276
...
275277

276278
@overload
277279
async def executemany_map(
278280
self,
279-
row_type: Type[GenericNewPyArrowTable],
281+
row_type: Type[GenericNewPanderaDataFrameModel],
280282
operation: str,
281283
seq_of_parameters: Sequence[Any],
282284
/,
283285
**options: Unpack[ExecuteOptions],
284-
) -> AsyncCursor[Never, Never, GenericNewPyArrowTable]:
286+
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
285287
...
286288

287289
@override
288290
async def executemany_map(
289291
self,
290292
row_type: Union[
291293
Type[GenericNewRowType],
292-
Type[GenericNewPanderaDataFrameModel],
293-
Type[GenericNewPandasDataFlame],
294+
Type[GenericNewPandasDataFrame],
294295
Type[GenericNewPyArrowTable],
296+
Type[GenericNewPanderaDataFrameModel],
295297
],
296298
operation: str,
297299
seq_of_parameters: Sequence[Any],

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

Lines changed: 27 additions & 25 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
@@ -139,44 +141,44 @@ async def execute_map(
139141
@overload
140142
async def execute_map(
141143
self,
142-
row_type: Type[GenericNewPanderaDataFrameModel],
144+
row_type: Type[GenericNewPandasDataFrame],
143145
operation: str,
144146
parameters: "Optional[Any]" = None,
145147
/,
146148
**options: Unpack[ExecuteOptions],
147-
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
149+
) -> "AsyncCursor[Never, GenericNewPandasDataFrame, Never]":
148150
...
149151

150152
@overload
151153
async def execute_map(
152154
self,
153-
row_type: Type[GenericNewPandasDataFlame],
155+
row_type: Type[GenericNewPyArrowTable],
154156
operation: str,
155157
parameters: "Optional[Any]" = None,
156158
/,
157159
**options: Unpack[ExecuteOptions],
158-
) -> "AsyncCursor[Never, GenericNewPandasDataFlame, Never]":
160+
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
159161
...
160162

161163
@overload
162164
async def execute_map(
163165
self,
164-
row_type: Type[GenericNewPyArrowTable],
166+
row_type: Type[GenericNewPanderaDataFrameModel],
165167
operation: str,
166168
parameters: "Optional[Any]" = None,
167169
/,
168170
**options: Unpack[ExecuteOptions],
169-
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
171+
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
170172
...
171173

172174
@override
173175
async def execute_map(
174176
self,
175177
row_type: Union[
176178
Type[GenericNewRowType],
177-
Type[GenericNewPanderaDataFrameModel],
178-
Type[GenericNewPandasDataFlame],
179+
Type[GenericNewPandasDataFrame],
179180
Type[GenericNewPyArrowTable],
181+
Type[GenericNewPanderaDataFrameModel],
180182
],
181183
operation: str,
182184
parameters: "Optional[Any]" = None,
@@ -215,44 +217,44 @@ async def executemany_map(
215217
@overload
216218
async def executemany_map(
217219
self,
218-
row_type: Type[GenericNewPanderaDataFrameModel],
220+
row_type: Type[GenericNewPandasDataFrame],
219221
operation: str,
220222
seq_of_parameters: "Sequence[Any]",
221223
/,
222224
**options: Unpack[ExecuteOptions],
223-
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
225+
) -> "AsyncCursor[Never, GenericNewPandasDataFrame, Never]":
224226
...
225227

226228
@overload
227229
async def executemany_map(
228230
self,
229-
row_type: Type[GenericNewPandasDataFlame],
231+
row_type: Type[GenericNewPyArrowTable],
230232
operation: str,
231233
seq_of_parameters: "Sequence[Any]",
232234
/,
233235
**options: Unpack[ExecuteOptions],
234-
) -> "AsyncCursor[Never, GenericNewPandasDataFlame, Never]":
236+
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
235237
...
236238

237239
@overload
238240
async def executemany_map(
239241
self,
240-
row_type: Type[GenericNewPyArrowTable],
242+
row_type: Type[GenericNewPanderaDataFrameModel],
241243
operation: str,
242244
seq_of_parameters: "Sequence[Any]",
243245
/,
244246
**options: Unpack[ExecuteOptions],
245-
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
247+
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
246248
...
247249

248250
@override
249251
async def executemany_map(
250252
self,
251253
row_type: Union[
252254
Type[GenericNewRowType],
253-
Type[GenericNewPanderaDataFrameModel],
254-
Type[GenericNewPandasDataFlame],
255+
Type[GenericNewPandasDataFrame],
255256
Type[GenericNewPyArrowTable],
257+
Type[GenericNewPanderaDataFrameModel],
256258
],
257259
operation: str,
258260
seq_of_parameters: "Sequence[Any]",
@@ -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

0 commit comments

Comments
 (0)