Merge pull request #62 from yassun7010/add_type_testcases

fix: execute() type hint.
yassun7010 · Feb 2, 2024 · c7c997a · c7c997a
2 parents 499598b + 4992944
commit c7c997a
Show file tree

Hide file tree

Showing 16 changed files with 470 additions and 195 deletions.
diff --git a/turu-snowflake/src/turu/snowflake/async_connection.py b/turu-snowflake/src/turu/snowflake/async_connection.py
@@ -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
 
@@ -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/),
@@ -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.
 
@@ -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,
@@ -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],

diff --git a/turu-snowflake/src/turu/snowflake/async_cursor.py b/turu-snowflake/src/turu/snowflake/async_cursor.py
@@ -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
 
@@ -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__(
@@ -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:
@@ -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.
 
@@ -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,
@@ -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]",
@@ -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),
  )