Skip to content

Commit d1d42cd

Browse files
ya7010ya7010
authored
Merge branch 'main' into dependabot/pip/pytest-8.0.0
2 parents fd9f526 + 2ad940d commit d1d42cd

16 files changed

+312
-250
lines changed

‎turu-snowflake/poetry.lock‎

Lines changed: 137 additions & 138 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎turu-snowflake/pyproject.toml‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,7 @@ classifiers = [
2828
]
2929

3030
[tool.poetry.extras]
31-
pandas = ["pandas"]
32-
pyarrow = ["pyarrow", "pyarrow-stubs"]
31+
pandas = ["pandas", "pyarrow", "pandas-stubs"]
3332
pandera = ["pandera"]
3433

3534
[tool.poetry.dependencies]

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

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -183,23 +183,23 @@ async def execute_map(
183183
@overload
184184
async def execute_map(
185185
self,
186-
row_type: Type[GenericNewPyArrowTable],
186+
row_type: Type[GenericNewPanderaDataFrameModel],
187187
operation: str,
188188
parameters: "Optional[Any]" = None,
189189
/,
190190
**options: Unpack[ExecuteOptions],
191-
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
191+
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
192192
...
193193

194194
@overload
195195
async def execute_map(
196196
self,
197-
row_type: Type[GenericNewPanderaDataFrameModel],
197+
row_type: Type[GenericNewPyArrowTable],
198198
operation: str,
199199
parameters: "Optional[Any]" = None,
200200
/,
201201
**options: Unpack[ExecuteOptions],
202-
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
202+
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
203203
...
204204

205205
@override
@@ -267,23 +267,23 @@ async def executemany_map(
267267
@overload
268268
async def executemany_map(
269269
self,
270-
row_type: Type[GenericNewPyArrowTable],
270+
row_type: Type[GenericNewPanderaDataFrameModel],
271271
operation: str,
272272
seq_of_parameters: Sequence[Any],
273273
/,
274274
**options: Unpack[ExecuteOptions],
275-
) -> AsyncCursor[Never, Never, GenericNewPyArrowTable]:
275+
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
276276
...
277277

278278
@overload
279279
async def executemany_map(
280280
self,
281-
row_type: Type[GenericNewPanderaDataFrameModel],
281+
row_type: Type[GenericNewPyArrowTable],
282282
operation: str,
283283
seq_of_parameters: Sequence[Any],
284284
/,
285285
**options: Unpack[ExecuteOptions],
286-
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
286+
) -> AsyncCursor[Never, Never, GenericNewPyArrowTable]:
287287
...
288288

289289
@override

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

Lines changed: 26 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import asyncio
22
from typing import (
33
Any,
4+
AsyncIterator,
45
Generic,
5-
Iterator,
66
List,
77
Optional,
88
Sequence,
@@ -17,18 +17,17 @@
1717
import turu.core.async_cursor
1818
import turu.core.cursor
1919
import turu.core.mock
20-
import turu.snowflake.record.async_record_cursor
2120
from turu.core.cursor import map_row
2221
from turu.snowflake.cursor import (
23-
GenericNewPandasDataFrame,
24-
GenericNewPyArrowTable,
2522
GenericNewRowType,
26-
GenericPandasDataFrame,
27-
GenericPyArrowTable,
2823
GenericRowType,
2924
)
3025
from turu.snowflake.features import (
26+
GenericNewPandasDataFrame,
3127
GenericNewPanderaDataFrameModel,
28+
GenericNewPyArrowTable,
29+
GenericPandasDataFrame,
30+
GenericPyArrowTable,
3231
PandasDataFrame,
3332
PanderaDataFrame,
3433
PanderaDataFrameModel,
@@ -152,23 +151,23 @@ async def execute_map(
152151
@overload
153152
async def execute_map(
154153
self,
155-
row_type: Type[GenericNewPyArrowTable],
154+
row_type: Type[GenericNewPanderaDataFrameModel],
156155
operation: str,
157156
parameters: "Optional[Any]" = None,
158157
/,
159158
**options: Unpack[ExecuteOptions],
160-
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
159+
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
161160
...
162161

163162
@overload
164163
async def execute_map(
165164
self,
166-
row_type: Type[GenericNewPanderaDataFrameModel],
165+
row_type: Type[GenericNewPyArrowTable],
167166
operation: str,
168167
parameters: "Optional[Any]" = None,
169168
/,
170169
**options: Unpack[ExecuteOptions],
171-
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
170+
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
172171
...
173172

174173
@override
@@ -228,23 +227,23 @@ async def executemany_map(
228227
@overload
229228
async def executemany_map(
230229
self,
231-
row_type: Type[GenericNewPyArrowTable],
230+
row_type: Type[GenericNewPanderaDataFrameModel],
232231
operation: str,
233232
seq_of_parameters: "Sequence[Any]",
234233
/,
235234
**options: Unpack[ExecuteOptions],
236-
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
235+
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
237236
...
238237

239238
@overload
240239
async def executemany_map(
241240
self,
242-
row_type: Type[GenericNewPanderaDataFrameModel],
241+
row_type: Type[GenericNewPyArrowTable],
243242
operation: str,
244243
seq_of_parameters: "Sequence[Any]",
245244
/,
246245
**options: Unpack[ExecuteOptions],
247-
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
246+
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
248247
...
249248

250249
@override
@@ -313,12 +312,11 @@ async def fetch_arrow_all(self) -> GenericPyArrowTable:
313312
self._raw_cursor.fetch_arrow_all(force_return_table=True),
314313
)
315314

316-
async def fetch_arrow_batches(self) -> Iterator[GenericPyArrowTable]:
315+
async def fetch_arrow_batches(self) -> AsyncIterator[GenericPyArrowTable]:
317316
"""Fetches Arrow Tables in batches, where 'batch' refers to Snowflake Chunk."""
318317

319-
return cast(
320-
Iterator[GenericPyArrowTable], self._raw_cursor.fetch_arrow_batches()
321-
)
318+
for batch in self._raw_cursor.fetch_arrow_batches():
319+
yield cast(GenericPyArrowTable, batch)
322320

323321
async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
324322
"""Fetch Pandas dataframes."""
@@ -330,13 +328,16 @@ async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
330328

331329
return cast(GenericPandasDataFrame, df)
332330

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

336-
return cast(
337-
Iterator[GenericPandasDataFrame],
338-
self._raw_cursor.fetch_pandas_batches(**kwargs),
339-
)
336+
for batch in self._raw_cursor.fetch_pandas_batches(**kwargs):
337+
yield cast(
338+
GenericPandasDataFrame,
339+
batch,
340+
)
340341

341342
@override
342343
async def __anext__(self) -> GenericRowType:
@@ -401,4 +402,6 @@ async def _execute_async(
401402
def _AsyncRecordCursor(
402403
self,
403404
):
405+
import turu.snowflake.record.async_record_cursor
406+
404407
return turu.snowflake.record.async_record_cursor.AsyncRecordCursor

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

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -193,23 +193,23 @@ def execute_map(
193193
@overload
194194
def execute_map(
195195
self,
196-
row_type: Type[GenericNewPyArrowTable],
196+
row_type: Type[GenericNewPanderaDataFrameModel],
197197
operation: str,
198198
parameters: "Optional[Any]" = None,
199199
/,
200200
**options: Unpack[ExecuteOptions],
201-
) -> Cursor[Never, Never, GenericNewPyArrowTable]:
201+
) -> Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]:
202202
...
203203

204204
@overload
205205
def execute_map(
206206
self,
207-
row_type: Type[GenericNewPanderaDataFrameModel],
207+
row_type: Type[GenericNewPyArrowTable],
208208
operation: str,
209209
parameters: "Optional[Any]" = None,
210210
/,
211211
**options: Unpack[ExecuteOptions],
212-
) -> Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]:
212+
) -> Cursor[Never, Never, GenericNewPyArrowTable]:
213213
...
214214

215215
@override
@@ -277,23 +277,23 @@ def executemany_map(
277277
@overload
278278
def executemany_map(
279279
self,
280-
row_type: Type[GenericNewPyArrowTable],
280+
row_type: Type[GenericNewPanderaDataFrameModel],
281281
operation: str,
282282
seq_of_parameters: Sequence[Any],
283283
/,
284284
**options: Unpack[ExecuteOptions],
285-
) -> Cursor[Never, Never, GenericNewPyArrowTable]:
285+
) -> Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]:
286286
...
287287

288288
@overload
289289
def executemany_map(
290290
self,
291-
row_type: Type[GenericNewPanderaDataFrameModel],
291+
row_type: Type[GenericNewPyArrowTable],
292292
operation: str,
293293
seq_of_parameters: Sequence[Any],
294294
/,
295295
**options: Unpack[ExecuteOptions],
296-
) -> Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]:
296+
) -> Cursor[Never, Never, GenericNewPyArrowTable]:
297297
...
298298

299299
@override

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

Lines changed: 12 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
Tuple,
99
Type,
1010
TypedDict,
11-
TypeVar,
1211
Union,
1312
cast,
1413
overload,
@@ -18,7 +17,11 @@
1817
import turu.core.mock
1918
from turu.core.cursor import GenericNewRowType, GenericRowType
2019
from turu.snowflake.features import (
20+
GenericNewPandasDataFrame,
2121
GenericNewPanderaDataFrameModel,
22+
GenericNewPyArrowTable,
23+
GenericPandasDataFrame,
24+
GenericPyArrowTable,
2225
PandasDataFrame,
2326
PanderaDataFrame,
2427
PanderaDataFrameModel,
@@ -37,13 +40,6 @@ class ExecuteOptions(TypedDict, total=False):
3740
"""number of statements"""
3841

3942

40-
GenericPyArrowTable = TypeVar("GenericPyArrowTable", bound=PyArrowTable)
41-
GenericNewPyArrowTable = TypeVar("GenericNewPyArrowTable", bound=PyArrowTable)
42-
43-
GenericPandasDataFrame = TypeVar("GenericPandasDataFrame", bound=PandasDataFrame)
44-
GenericNewPandasDataFrame = TypeVar("GenericNewPandasDataFrame", bound=PandasDataFrame)
45-
46-
4743
class Cursor(
4844
Generic[GenericRowType, GenericPandasDataFrame, GenericPyArrowTable],
4945
turu.core.cursor.Cursor[GenericRowType, Any],
@@ -229,24 +225,24 @@ def executemany_map(
229225
@overload
230226
def executemany_map(
231227
self,
232-
row_type: Type[GenericNewPyArrowTable],
228+
row_type: Type[GenericNewPanderaDataFrameModel],
233229
operation: str,
234-
seq_of_parameters: "Sequence[Any]",
230+
seq_of_parameters: Sequence[Any],
235231
/,
236232
**options: Unpack[ExecuteOptions],
237-
) -> "Cursor[Never, Never, GenericNewPyArrowTable]":
238-
pass
233+
) -> "Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
234+
...
239235

240236
@overload
241237
def executemany_map(
242238
self,
243-
row_type: Type[GenericNewPanderaDataFrameModel],
239+
row_type: Type[GenericNewPyArrowTable],
244240
operation: str,
245-
seq_of_parameters: Sequence[Any],
241+
seq_of_parameters: "Sequence[Any]",
246242
/,
247243
**options: Unpack[ExecuteOptions],
248-
) -> "Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
249-
...
244+
) -> "Cursor[Never, Never, GenericNewPyArrowTable]":
245+
pass
250246

251247
@override
252248
def executemany_map(

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,22 +15,34 @@ class _NotSupportFeatureT(Generic[T]):
1515

1616
USE_PANDAS = True
1717
PandasDataFrame: TypeAlias = pandas.DataFrame # type: ignore
18+
GenericPandasDataFrame = TypeVar("GenericPandasDataFrame", bound=PandasDataFrame)
19+
GenericNewPandasDataFrame = TypeVar(
20+
"GenericNewPandasDataFrame", bound=PandasDataFrame
21+
)
1822

1923
except ImportError:
2024
USE_PANDAS = False
2125
PandasDataFrame: TypeAlias = _NotSupportFeature # type: ignore
26+
GenericPandasDataFrame = TypeVar("GenericPandasDataFrame", bound=_NotSupportFeature)
27+
GenericNewPandasDataFrame = TypeVar(
28+
"GenericNewPandasDataFrame", bound=_NotSupportFeature
29+
)
2230

2331

2432
try:
2533
import pyarrow # type: ignore[import] # noqa: F401
2634

2735
USE_PYARROW = True
2836
PyArrowTable: TypeAlias = pyarrow.Table # type: ignore
37+
GenericPyArrowTable = TypeVar("GenericPyArrowTable", bound=PyArrowTable)
38+
GenericNewPyArrowTable = TypeVar("GenericNewPyArrowTable", bound=PyArrowTable)
2939

3040

3141
except ImportError:
3242
USE_PYARROW = False
3343
PyArrowTable: TypeAlias = _NotSupportFeature # type: ignore
44+
GenericPyArrowTable = TypeVar("GenericPyArrowTable", bound=_NotSupportFeature)
45+
GenericNewPyArrowTable = TypeVar("GenericNewPyArrowTable", bound=_NotSupportFeature)
3446

3547
try:
3648
from typing import TypeVar

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,12 @@
1010
from turu.core.cursor import GenericRowType
1111
from turu.core.mock.connection import CSVOptions
1212
from turu.core.mock.exception import TuruCsvHeaderOptionRequiredError
13-
from turu.snowflake.cursor import GenericPandasDataFrame, GenericPyArrowTable
1413
from turu.snowflake.features import (
1514
USE_PANDAS,
1615
USE_PYARROW,
16+
GenericPandasDataFrame,
1717
GenericPanderaDataFrameModel,
18+
GenericPyArrowTable,
1819
PandasDataFrame,
1920
PanderaDataFrameModel,
2021
PyArrowTable,

0 commit comments

Comments
 (0)