Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: execute() type hint. #62

Merged
merged 4 commits into from
Feb 2, 2024
Merged

fix: execute() type hint. #62

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c5e5fe5
test: add cursor type tests.
yassun7010 Feb 2, 2024
b5f6083
fix: type hint.
yassun7010 Feb 2, 2024
fcf07ce
fix: type.
yassun7010 Feb 2, 2024
4992944
Merge branch 'main' into add_type_testcases
yassun7010 Feb 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 21 additions & 19 deletions turu-snowflake/src/turu/snowflake/async_connection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,12 @@
import turu.core.cursor
import turu.core.mock
from turu.core.cursor import GenericNewRowType
from turu.snowflake.cursor import GenericNewPandasDataFlame, GenericNewPyArrowTable
from turu.snowflake.cursor import GenericNewPandasDataFrame, GenericNewPyArrowTable
from turu.snowflake.features import (
GenericNewPanderaDataFrameModel,
PandasDataFrame,
PanderaDataFrame,
PyArrowTable,
)
from typing_extensions import Never, Self, Unpack, override

Expand Down Expand Up @@ -112,7 +114,7 @@ async def execute(
parameters: Optional[Any] = None,
/,
**options: Unpack[ExecuteOptions],
) -> AsyncCursor[Tuple[Any], Never, Never]:
) -> AsyncCursor[Tuple[Any], PandasDataFrame, PyArrowTable]:
"""Prepare and execute a database operation (query or command).

This is not defined in [PEP 249](https://peps.python.org/pep-0249/),
Expand All @@ -136,7 +138,7 @@ async def executemany(
seq_of_parameters: Sequence[Any],
/,
**options: Unpack[ExecuteOptions],
) -> AsyncCursor[Tuple[Any], Never, Never]:
) -> AsyncCursor[Tuple[Any], PandasDataFrame, PyArrowTable]:
"""Prepare a database operation (query or command)
and then execute it against all parameter sequences or mappings.

Expand Down Expand Up @@ -170,44 +172,44 @@ async def execute_map(
@overload
async def execute_map(
self,
row_type: Type[GenericNewPandasDataFlame],
row_type: Type[GenericNewPandasDataFrame],
operation: str,
parameters: "Optional[Any]" = None,
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Never, GenericNewPandasDataFlame, Never]":
) -> "AsyncCursor[Never, GenericNewPandasDataFrame, Never]":
...

@overload
async def execute_map(
self,
row_type: Type[GenericNewPanderaDataFrameModel],
row_type: Type[GenericNewPyArrowTable],
operation: str,
parameters: "Optional[Any]" = None,
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
...

@overload
async def execute_map(
self,
row_type: Type[GenericNewPyArrowTable],
row_type: Type[GenericNewPanderaDataFrameModel],
operation: str,
parameters: "Optional[Any]" = None,
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
...

@override
async def execute_map(
self,
row_type: Union[
Type[GenericNewRowType],
Type[GenericNewPanderaDataFrameModel],
Type[GenericNewPandasDataFlame],
Type[GenericNewPandasDataFrame],
Type[GenericNewPyArrowTable],
Type[GenericNewPanderaDataFrameModel],
],
operation: str,
parameters: Optional[Any] = None,
Expand Down Expand Up @@ -254,44 +256,44 @@ async def executemany_map(
@overload
async def executemany_map(
self,
row_type: Type[GenericNewPanderaDataFrameModel],
row_type: Type[GenericNewPandasDataFrame],
operation: str,
seq_of_parameters: Sequence[Any],
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
) -> AsyncCursor[Never, GenericNewPandasDataFrame, Never]:
...

@overload
async def executemany_map(
self,
row_type: Type[GenericNewPandasDataFlame],
row_type: Type[GenericNewPyArrowTable],
operation: str,
seq_of_parameters: Sequence[Any],
/,
**options: Unpack[ExecuteOptions],
) -> AsyncCursor[Never, GenericNewPandasDataFlame, Never]:
) -> AsyncCursor[Never, Never, GenericNewPyArrowTable]:
...

@overload
async def executemany_map(
self,
row_type: Type[GenericNewPyArrowTable],
row_type: Type[GenericNewPanderaDataFrameModel],
operation: str,
seq_of_parameters: Sequence[Any],
/,
**options: Unpack[ExecuteOptions],
) -> AsyncCursor[Never, Never, GenericNewPyArrowTable]:
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
...

@override
async def executemany_map(
self,
row_type: Union[
Type[GenericNewRowType],
Type[GenericNewPanderaDataFrameModel],
Type[GenericNewPandasDataFlame],
Type[GenericNewPandasDataFrame],
Type[GenericNewPyArrowTable],
Type[GenericNewPanderaDataFrameModel],
],
operation: str,
seq_of_parameters: Sequence[Any],
Expand Down
52 changes: 27 additions & 25 deletions turu-snowflake/src/turu/snowflake/async_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,17 +20,19 @@
import turu.snowflake.record.async_record_cursor
from turu.core.cursor import map_row
from turu.snowflake.cursor import (
GenericNewPandasDataFlame,
GenericNewPandasDataFrame,
GenericNewPyArrowTable,
GenericNewRowType,
GenericPandasDataFlame,
GenericPandasDataFrame,
GenericPyArrowTable,
GenericRowType,
)
from turu.snowflake.features import (
GenericNewPanderaDataFrameModel,
PandasDataFrame,
PanderaDataFrame,
PanderaDataFrameModel,
PyArrowTable,
)
from typing_extensions import Never, Self, Unpack, override

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


class AsyncCursor(
Generic[GenericRowType, GenericPandasDataFlame, GenericPyArrowTable],
Generic[GenericRowType, GenericPandasDataFrame, GenericPyArrowTable],
turu.core.async_cursor.AsyncCursor[GenericRowType, Any],
):
def __init__(
Expand Down Expand Up @@ -80,7 +82,7 @@ async def execute(
parameters: Optional[Any] = None,
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Tuple[Any], Never, Never]":
) -> "AsyncCursor[Tuple[Any], PandasDataFrame, PyArrowTable]":
"""Prepare and execute a database operation (query or command).

Parameters:
Expand All @@ -104,7 +106,7 @@ async def executemany(
seq_of_parameters: Sequence[Any],
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Tuple[Any], Never, Never]":
) -> "AsyncCursor[Tuple[Any], PandasDataFrame, PyArrowTable]":
"""Prepare a database operation (query or command)
and then execute it against all parameter sequences or mappings.

Expand Down Expand Up @@ -139,44 +141,44 @@ async def execute_map(
@overload
async def execute_map(
self,
row_type: Type[GenericNewPanderaDataFrameModel],
row_type: Type[GenericNewPandasDataFrame],
operation: str,
parameters: "Optional[Any]" = None,
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
) -> "AsyncCursor[Never, GenericNewPandasDataFrame, Never]":
...

@overload
async def execute_map(
self,
row_type: Type[GenericNewPandasDataFlame],
row_type: Type[GenericNewPyArrowTable],
operation: str,
parameters: "Optional[Any]" = None,
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Never, GenericNewPandasDataFlame, Never]":
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
...

@overload
async def execute_map(
self,
row_type: Type[GenericNewPyArrowTable],
row_type: Type[GenericNewPanderaDataFrameModel],
operation: str,
parameters: "Optional[Any]" = None,
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
...

@override
async def execute_map(
self,
row_type: Union[
Type[GenericNewRowType],
Type[GenericNewPanderaDataFrameModel],
Type[GenericNewPandasDataFlame],
Type[GenericNewPandasDataFrame],
Type[GenericNewPyArrowTable],
Type[GenericNewPanderaDataFrameModel],
],
operation: str,
parameters: "Optional[Any]" = None,
Expand Down Expand Up @@ -215,44 +217,44 @@ async def executemany_map(
@overload
async def executemany_map(
self,
row_type: Type[GenericNewPanderaDataFrameModel],
row_type: Type[GenericNewPandasDataFrame],
operation: str,
seq_of_parameters: "Sequence[Any]",
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
) -> "AsyncCursor[Never, GenericNewPandasDataFrame, Never]":
...

@overload
async def executemany_map(
self,
row_type: Type[GenericNewPandasDataFlame],
row_type: Type[GenericNewPyArrowTable],
operation: str,
seq_of_parameters: "Sequence[Any]",
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Never, GenericNewPandasDataFlame, Never]":
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
...

@overload
async def executemany_map(
self,
row_type: Type[GenericNewPyArrowTable],
row_type: Type[GenericNewPanderaDataFrameModel],
operation: str,
seq_of_parameters: "Sequence[Any]",
/,
**options: Unpack[ExecuteOptions],
) -> "AsyncCursor[Never, Never, GenericNewPyArrowTable]":
) -> "AsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]":
...

@override
async def executemany_map(
self,
row_type: Union[
Type[GenericNewRowType],
Type[GenericNewPanderaDataFrameModel],
Type[GenericNewPandasDataFlame],
Type[GenericNewPandasDataFrame],
Type[GenericNewPyArrowTable],
Type[GenericNewPanderaDataFrameModel],
],
operation: str,
seq_of_parameters: "Sequence[Any]",
Expand Down Expand Up @@ -318,21 +320,21 @@ async def fetch_arrow_batches(self) -> Iterator[GenericPyArrowTable]:
Iterator[GenericPyArrowTable], self._raw_cursor.fetch_arrow_batches()
)

async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFlame:
async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame:
"""Fetch Pandas dataframes."""

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

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

return cast(GenericPandasDataFlame, df)
return cast(GenericPandasDataFrame, df)

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

return cast(
Iterator[GenericPandasDataFlame],
Iterator[GenericPandasDataFrame],
self._raw_cursor.fetch_pandas_batches(**kwargs),
)

Expand Down
Loading
Loading