diff --git a/turu-snowflake/src/turu/snowflake/async_connection.py b/turu-snowflake/src/turu/snowflake/async_connection.py index 9a9445c..0a1fc57 100644 --- 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,34 +172,34 @@ 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 @@ -205,9 +207,9 @@ 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,34 +256,34 @@ 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 @@ -289,9 +291,9 @@ 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 index ee3bb2c..43f09a9 100644 --- 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,34 +141,34 @@ 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 @@ -174,9 +176,9 @@ 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,34 +217,34 @@ 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 @@ -250,9 +252,9 @@ 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,7 +320,7 @@ 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) @@ -326,13 +328,13 @@ async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFlame: 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), ) diff --git a/turu-snowflake/src/turu/snowflake/connection.py b/turu-snowflake/src/turu/snowflake/connection.py index ce1adce..8453bdb 100644 --- a/turu-snowflake/src/turu/snowflake/connection.py +++ b/turu-snowflake/src/turu/snowflake/connection.py @@ -9,7 +9,9 @@ from turu.core.cursor import GenericNewRowType from turu.snowflake.features import ( GenericNewPanderaDataFrameModel, + PandasDataFrame, PanderaDataFrame, + PyArrowTable, ) from typing_extensions import Never, Unpack, override @@ -18,7 +20,7 @@ from .cursor import ( Cursor, ExecuteOptions, - GenericNewPandasDataFlame, + GenericNewPandasDataFrame, GenericNewPyArrowTable, Self, ) @@ -118,7 +120,7 @@ def execute( parameters: Optional[Any] = None, /, **options: Unpack[ExecuteOptions], - ) -> Cursor[Tuple[Any], Never, Never]: + ) -> Cursor[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/), @@ -142,7 +144,7 @@ def executemany( seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> Cursor[Tuple[Any], Never, Never]: + ) -> Cursor[Tuple[Any], PandasDataFrame, PyArrowTable]: """Prepare a database operation (query or command) and then execute it against all parameter sequences or mappings. @@ -174,34 +176,34 @@ def execute_map( @overload def execute_map( self, - row_type: Type[GenericNewPanderaDataFrameModel], + row_type: Type[GenericNewPandasDataFrame], operation: str, parameters: "Optional[Any]" = None, /, **options: Unpack[ExecuteOptions], - ) -> Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]: + ) -> Cursor[Never, GenericNewPandasDataFrame, Never]: ... @overload def execute_map( self, - row_type: Type[GenericNewPandasDataFlame], + row_type: Type[GenericNewPyArrowTable], operation: str, parameters: "Optional[Any]" = None, /, **options: Unpack[ExecuteOptions], - ) -> Cursor[Never, GenericNewPandasDataFlame, Never]: + ) -> Cursor[Never, Never, GenericNewPyArrowTable]: ... @overload def execute_map( self, - row_type: Type[GenericNewPyArrowTable], + row_type: Type[GenericNewPanderaDataFrameModel], operation: str, parameters: "Optional[Any]" = None, /, **options: Unpack[ExecuteOptions], - ) -> Cursor[Never, Never, GenericNewPyArrowTable]: + ) -> Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]: ... @override @@ -209,9 +211,9 @@ 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, @@ -258,34 +260,34 @@ def executemany_map( @overload def executemany_map( self, - row_type: Type[GenericNewPanderaDataFrameModel], + row_type: Type[GenericNewPandasDataFrame], operation: str, seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]: + ) -> Cursor[Never, GenericNewPandasDataFrame, Never]: ... @overload def executemany_map( self, - row_type: Type[GenericNewPandasDataFlame], + row_type: Type[GenericNewPyArrowTable], operation: str, seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> Cursor[Never, GenericNewPandasDataFlame, Never]: + ) -> Cursor[Never, Never, GenericNewPyArrowTable]: ... @overload def executemany_map( self, - row_type: Type[GenericNewPyArrowTable], + row_type: Type[GenericNewPanderaDataFrameModel], operation: str, seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> Cursor[Never, Never, GenericNewPyArrowTable]: + ) -> Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]: ... @override @@ -293,9 +295,9 @@ 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/cursor.py b/turu-snowflake/src/turu/snowflake/cursor.py index c0170a9..4b79d40 100644 --- a/turu-snowflake/src/turu/snowflake/cursor.py +++ b/turu-snowflake/src/turu/snowflake/cursor.py @@ -19,7 +19,7 @@ from turu.core.cursor import GenericNewRowType, GenericRowType from turu.snowflake.features import ( GenericNewPanderaDataFrameModel, - PandasDataFlame, + PandasDataFrame, PanderaDataFrame, PanderaDataFrameModel, PyArrowTable, @@ -40,12 +40,12 @@ class ExecuteOptions(TypedDict, total=False): GenericPyArrowTable = TypeVar("GenericPyArrowTable", bound=PyArrowTable) GenericNewPyArrowTable = TypeVar("GenericNewPyArrowTable", bound=PyArrowTable) -GenericPandasDataFlame = TypeVar("GenericPandasDataFlame", bound=PandasDataFlame) -GenericNewPandasDataFlame = TypeVar("GenericNewPandasDataFlame", bound=PandasDataFlame) +GenericPandasDataFrame = TypeVar("GenericPandasDataFrame", bound=PandasDataFrame) +GenericNewPandasDataFrame = TypeVar("GenericNewPandasDataFrame", bound=PandasDataFrame) class Cursor( - Generic[GenericRowType, GenericPandasDataFlame, GenericPyArrowTable], + Generic[GenericRowType, GenericPandasDataFrame, GenericPyArrowTable], turu.core.cursor.Cursor[GenericRowType, Any], ): def __init__( @@ -80,7 +80,7 @@ def execute( parameters: Optional[Any] = None, /, **options: Unpack[ExecuteOptions], - ) -> "Cursor[Tuple[Any], Never, Never]": + ) -> "Cursor[Tuple[Any], PandasDataFrame, PyArrowTable]": """Prepare and execute a database operation (query or command). Parameters: @@ -104,7 +104,7 @@ def executemany( seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> "Cursor[Tuple[Any], Never, Never]": + ) -> "Cursor[Tuple[Any], PandasDataFrame, PyArrowTable]": """Prepare a database operation (query or command) and then execute it against all parameter sequences or mappings. @@ -136,34 +136,34 @@ def execute_map( @overload def execute_map( self, - row_type: Type[GenericNewPanderaDataFrameModel], + row_type: Type[GenericNewPandasDataFrame], operation: str, parameters: "Optional[Any]" = None, /, **options: Unpack[ExecuteOptions], - ) -> "Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": + ) -> "Cursor[Never, GenericNewPandasDataFrame, Never]": ... @overload def execute_map( self, - row_type: Type[GenericNewPandasDataFlame], + row_type: Type[GenericNewPyArrowTable], operation: str, parameters: "Optional[Any]" = None, /, **options: Unpack[ExecuteOptions], - ) -> "Cursor[Never, GenericNewPandasDataFlame, Never]": + ) -> "Cursor[Never, Never, GenericNewPyArrowTable]": ... @overload def execute_map( self, - row_type: Type[GenericNewPyArrowTable], + row_type: Type[GenericNewPanderaDataFrameModel], operation: str, parameters: "Optional[Any]" = None, /, **options: Unpack[ExecuteOptions], - ) -> "Cursor[Never, Never, GenericNewPyArrowTable]": + ) -> "Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": ... @override @@ -171,9 +171,9 @@ 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, @@ -212,44 +212,44 @@ def executemany_map( @overload def executemany_map( self, - row_type: Type[GenericNewPanderaDataFrameModel], + row_type: Type[GenericNewPandasDataFrame], operation: str, - seq_of_parameters: Sequence[Any], + seq_of_parameters: "Sequence[Any]", /, **options: Unpack[ExecuteOptions], - ) -> "Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": - ... + ) -> "Cursor[Never, GenericNewPandasDataFrame, Never]": + pass @overload def executemany_map( self, - row_type: Type[GenericNewPandasDataFlame], + row_type: Type[GenericNewPyArrowTable], operation: str, seq_of_parameters: "Sequence[Any]", /, **options: Unpack[ExecuteOptions], - ) -> "Cursor[Never, GenericNewPandasDataFlame, Never]": + ) -> "Cursor[Never, Never, GenericNewPyArrowTable]": pass @overload def executemany_map( self, - row_type: Type[GenericNewPyArrowTable], + row_type: Type[GenericNewPanderaDataFrameModel], operation: str, - seq_of_parameters: "Sequence[Any]", + seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> "Cursor[Never, Never, GenericNewPyArrowTable]": - pass + ) -> "Cursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": + ... @override 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]", @@ -329,20 +329,20 @@ def fetch_arrow_batches(self) -> "Iterator[GenericPyArrowTable]": Iterator[GenericPyArrowTable], self._raw_cursor.fetch_arrow_batches() ) - def fetch_pandas_all(self, **kwargs) -> "GenericPandasDataFlame": + def fetch_pandas_all(self, **kwargs) -> "GenericPandasDataFrame": """Fetch a single Pandas dataframe.""" 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[assignment] - return cast(GenericPandasDataFlame, df) + return cast(GenericPandasDataFrame, df) - def fetch_pandas_batches(self, **kwargs) -> "Iterator[GenericPandasDataFlame]": + 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), ) diff --git a/turu-snowflake/src/turu/snowflake/features.py b/turu-snowflake/src/turu/snowflake/features.py index 2c71286..8f855c7 100644 --- a/turu-snowflake/src/turu/snowflake/features.py +++ b/turu-snowflake/src/turu/snowflake/features.py @@ -14,18 +14,18 @@ class _NotSupportFeatureT(Generic[T]): import pandas # type: ignore[import] # noqa: F401 USE_PANDAS = True - PandasDataFlame: TypeAlias = pandas.DataFrame # type: ignore + PandasDataFrame: TypeAlias = pandas.DataFrame # type: ignore except ImportError: USE_PANDAS = False - PandasDataFlame: TypeAlias = _NotSupportFeature # type: ignore + PandasDataFrame: TypeAlias = _NotSupportFeature # type: ignore try: import pyarrow # type: ignore[import] # noqa: F401 USE_PYARROW = True - PyArrowTable = pyarrow.Table # type: ignore + PyArrowTable: TypeAlias = pyarrow.Table # type: ignore except ImportError: diff --git a/turu-snowflake/src/turu/snowflake/mock_async_connection.py b/turu-snowflake/src/turu/snowflake/mock_async_connection.py index 9ec12ee..f5bd570 100644 --- a/turu-snowflake/src/turu/snowflake/mock_async_connection.py +++ b/turu-snowflake/src/turu/snowflake/mock_async_connection.py @@ -10,10 +10,10 @@ from turu.core.cursor import GenericRowType from turu.core.mock.connection import CSVOptions from turu.core.mock.exception import TuruCsvHeaderOptionRequiredError -from turu.snowflake.cursor import GenericPandasDataFlame, GenericPyArrowTable +from turu.snowflake.cursor import GenericPandasDataFrame, GenericPyArrowTable from turu.snowflake.features import ( GenericPanderaDataFrameModel, - PandasDataFlame, + PandasDataFrame, PanderaDataFrameModel, PyArrowTable, ) @@ -52,9 +52,9 @@ def inject_response( @overload def inject_response( self, - row_type: Type[GenericPanderaDataFrameModel], + row_type: Type[GenericPandasDataFrame], response: Union[ - Sequence[GenericPandasDataFlame], GenericPandasDataFlame, Exception + Sequence[GenericPandasDataFrame], GenericPandasDataFrame, Exception ], ) -> Self: ... @@ -62,18 +62,18 @@ def inject_response( @overload def inject_response( self, - row_type: Type[GenericPandasDataFlame], - response: Union[ - Sequence[GenericPandasDataFlame], GenericPandasDataFlame, Exception - ], + row_type: Type[GenericPyArrowTable], + response: Union[Sequence[GenericPyArrowTable], GenericPyArrowTable, Exception], ) -> Self: ... @overload def inject_response( self, - row_type: Type[GenericPyArrowTable], - response: Union[Sequence[GenericPyArrowTable], GenericPyArrowTable, Exception], + row_type: Type[GenericPanderaDataFrameModel], + response: Union[ + Sequence[GenericPandasDataFrame], GenericPandasDataFrame, Exception + ], ) -> Self: ... @@ -82,14 +82,14 @@ def inject_response( # type: ignore[override] self, row_type: Union[ Type[GenericRowType], - Type[GenericPanderaDataFrameModel], - Type[GenericPandasDataFlame], + Type[GenericPandasDataFrame], Type[GenericPyArrowTable], + Type[GenericPanderaDataFrameModel], None, ], response: Union[Sequence[Any], Any, Exception] = None, ) -> Self: - if row_type is not None and isinstance(response, PandasDataFlame): + if row_type is not None and isinstance(response, PandasDataFrame): response = (response,) self._turu_mock_store.inject_response( @@ -103,15 +103,15 @@ def inject_response_from_csv( # type: ignore[override] self, row_type: Union[ Type[GenericRowType], - Type[GenericPanderaDataFrameModel], - Type[GenericPandasDataFlame], + Type[GenericPandasDataFrame], Type[GenericPyArrowTable], + Type[GenericPanderaDataFrameModel], ], filepath: Union[str, pathlib.Path], **options: Unpack[CSVOptions], ) -> Self: if row_type is not None: - if issubclass(row_type, (PandasDataFlame, PanderaDataFrameModel)): + if issubclass(row_type, (PandasDataFrame, PanderaDataFrameModel)): import pandas if options.get("header", True) is False: diff --git a/turu-snowflake/src/turu/snowflake/mock_async_cursor.py b/turu-snowflake/src/turu/snowflake/mock_async_cursor.py index 61dec30..7b75beb 100644 --- a/turu-snowflake/src/turu/snowflake/mock_async_cursor.py +++ b/turu-snowflake/src/turu/snowflake/mock_async_cursor.py @@ -14,15 +14,17 @@ import turu.core.mock from turu.core.cursor import GenericNewRowType, GenericRowType from turu.snowflake.cursor import ( - GenericNewPandasDataFlame, + GenericNewPandasDataFrame, GenericNewPyArrowTable, - GenericPandasDataFlame, + GenericPandasDataFrame, GenericPyArrowTable, ) from turu.snowflake.features import ( GenericNewPanderaDataFrameModel, + PandasDataFrame, PanderaDataFrame, PanderaDataFrameModel, + PyArrowTable, ) from typing_extensions import Never, Self, Unpack, override @@ -31,7 +33,7 @@ class MockAsyncCursor( # type: ignore turu.core.mock.MockAsyncCursor[GenericRowType, Any], # type: ignore - AsyncCursor[GenericRowType, GenericPandasDataFlame, GenericPyArrowTable], # type: ignore + AsyncCursor[GenericRowType, GenericPandasDataFrame, GenericPyArrowTable], # type: ignore ): @override async def execute( @@ -40,7 +42,7 @@ async def execute( parameters: Optional[Any] = None, /, **options: Unpack[ExecuteOptions], - ) -> "MockAsyncCursor[Tuple[Any], Never, Never]": + ) -> "MockAsyncCursor[Tuple[Any], PandasDataFrame, PyArrowTable]": return cast(MockAsyncCursor, await super().execute(operation, parameters)) @override @@ -50,7 +52,7 @@ async def executemany( seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> "MockAsyncCursor[Tuple[Any], Never, Never]": + ) -> "MockAsyncCursor[Tuple[Any], PandasDataFrame, PyArrowTable]": return cast( MockAsyncCursor, await super().executemany(operation, seq_of_parameters) ) @@ -69,34 +71,34 @@ 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], - ) -> "MockAsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": + ) -> "MockAsyncCursor[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], - ) -> "MockAsyncCursor[Never, GenericNewPandasDataFlame, Never]": + ) -> "MockAsyncCursor[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], - ) -> "MockAsyncCursor[Never, Never, GenericNewPyArrowTable]": + ) -> "MockAsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": ... @override @@ -104,9 +106,9 @@ 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, @@ -134,34 +136,34 @@ 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], - ) -> "MockAsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": + ) -> "MockAsyncCursor[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], - ) -> "MockAsyncCursor[Never, GenericNewPandasDataFlame, Never]": + ) -> "MockAsyncCursor[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], - ) -> "MockAsyncCursor[Never, Never, GenericNewPyArrowTable]": + ) -> "MockAsyncCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": ... @override @@ -169,9 +171,9 @@ 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], @@ -205,7 +207,7 @@ async def fetch_arrow_all(self) -> GenericPyArrowTable: async def fetch_arrow_batches(self) -> Iterator[GenericPyArrowTable]: return iter([await self.fetch_arrow_all()]) - async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFlame: + async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame: df = await self.fetchone() if ( @@ -215,9 +217,9 @@ async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFlame: ): df = self._row_type.validate(df) # type: ignore - 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]: """Fetches a single Arrow Table.""" return iter([await self.fetch_pandas_all(**kwargs)]) diff --git a/turu-snowflake/src/turu/snowflake/mock_connection.py b/turu-snowflake/src/turu/snowflake/mock_connection.py index d56110f..22995a0 100644 --- a/turu-snowflake/src/turu/snowflake/mock_connection.py +++ b/turu-snowflake/src/turu/snowflake/mock_connection.py @@ -9,10 +9,10 @@ from turu.core.cursor import GenericRowType from turu.core.mock.connection import CSVOptions from turu.core.mock.exception import TuruCsvHeaderOptionRequiredError -from turu.snowflake.cursor import GenericPandasDataFlame, GenericPyArrowTable +from turu.snowflake.cursor import GenericPandasDataFrame, GenericPyArrowTable from turu.snowflake.features import ( GenericPanderaDataFrameModel, - PandasDataFlame, + PandasDataFrame, PanderaDataFrameModel, PyArrowTable, ) @@ -26,7 +26,9 @@ def __init__(self, *args, **kwargs): turu.core.mock.MockConnection.__init__(self) @override - def cursor(self) -> "turu.snowflake.mock_cursor.MockCursor[Never, Never, Never]": + def cursor( + self, + ) -> "turu.snowflake.mock_cursor.MockCursor[Never, Never, Never]": return turu.snowflake.mock_cursor.MockCursor(self._turu_mock_store) @overload @@ -50,9 +52,9 @@ def inject_response( @overload def inject_response( self, - row_type: Type[GenericPanderaDataFrameModel], + row_type: Type[GenericPandasDataFrame], response: Union[ - Sequence[GenericPandasDataFlame], GenericPandasDataFlame, Exception + Sequence[GenericPandasDataFrame], GenericPandasDataFrame, Exception ], ) -> Self: ... @@ -60,18 +62,18 @@ def inject_response( @overload def inject_response( self, - row_type: Type[GenericPandasDataFlame], - response: Union[ - Sequence[GenericPandasDataFlame], GenericPandasDataFlame, Exception - ], + row_type: Type[GenericPyArrowTable], + response: Union[Sequence[GenericPyArrowTable], GenericPyArrowTable, Exception], ) -> Self: ... @overload def inject_response( self, - row_type: Type[GenericPyArrowTable], - response: Union[Sequence[GenericPyArrowTable], GenericPyArrowTable, Exception], + row_type: Type[GenericPanderaDataFrameModel], + response: Union[ + Sequence[GenericPandasDataFrame], GenericPandasDataFrame, Exception + ], ) -> Self: ... @@ -80,14 +82,14 @@ def inject_response( # type: ignore[override] self, row_type: Union[ Type[GenericRowType], - Type[GenericPanderaDataFrameModel], - Type[GenericPandasDataFlame], + Type[GenericPandasDataFrame], Type[GenericPyArrowTable], + Type[GenericPanderaDataFrameModel], None, ], response: Union[Sequence[Any], Any, Exception] = None, ) -> Self: - if row_type is not None and isinstance(response, PandasDataFlame): + if row_type is not None and isinstance(response, PandasDataFrame): response = (response,) self._turu_mock_store.inject_response( @@ -101,15 +103,15 @@ def inject_response_from_csv( # type: ignore[override] self, row_type: Union[ Type[GenericRowType], - Type[GenericPanderaDataFrameModel], - Type[GenericPandasDataFlame], + Type[GenericPandasDataFrame], Type[GenericPyArrowTable], + Type[GenericPanderaDataFrameModel], ], filepath: Union[str, pathlib.Path], **options: Unpack[CSVOptions], ) -> Self: if row_type is not None: - if issubclass(row_type, (PandasDataFlame, PanderaDataFrameModel)): + if issubclass(row_type, (PandasDataFrame, PanderaDataFrameModel)): import pandas pd_options = {} diff --git a/turu-snowflake/src/turu/snowflake/mock_cursor.py b/turu-snowflake/src/turu/snowflake/mock_cursor.py index c364353..06b4cbf 100644 --- a/turu-snowflake/src/turu/snowflake/mock_cursor.py +++ b/turu-snowflake/src/turu/snowflake/mock_cursor.py @@ -15,17 +15,19 @@ from turu.core.cursor import GenericNewRowType from turu.snowflake.features import ( GenericNewPanderaDataFrameModel, + PandasDataFrame, PanderaDataFrame, PanderaDataFrameModel, + PyArrowTable, ) from typing_extensions import Never, Self, Unpack, override from .cursor import ( Cursor, ExecuteOptions, - GenericNewPandasDataFlame, + GenericNewPandasDataFrame, GenericNewPyArrowTable, - GenericPandasDataFlame, + GenericPandasDataFrame, GenericPyArrowTable, ) @@ -33,7 +35,7 @@ class MockCursor( # type: ignore turu.core.mock.MockCursor[turu.core.cursor.GenericRowType, Any], # type: ignore Cursor[ - turu.core.cursor.GenericRowType, GenericPandasDataFlame, GenericPyArrowTable + turu.core.cursor.GenericRowType, GenericPandasDataFrame, GenericPyArrowTable ], # type: ignore ): @override @@ -43,7 +45,7 @@ def execute( parameters: Optional[Any] = None, /, **options: Unpack[ExecuteOptions], - ) -> "MockCursor[Tuple[Any], Never, Never]": + ) -> "MockCursor[Tuple[Any], PandasDataFrame, PyArrowTable]": return cast(MockCursor, super().execute(operation, parameters)) @override @@ -53,7 +55,7 @@ def executemany( seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> "MockCursor[Tuple[Any], Never, Never]": + ) -> "MockCursor[Tuple[Any], PandasDataFrame, PyArrowTable]": return cast(MockCursor, super().executemany(operation, seq_of_parameters)) @overload @@ -70,34 +72,34 @@ def execute_map( @overload def execute_map( self, - row_type: Type[GenericNewPanderaDataFrameModel], + row_type: Type[GenericNewPandasDataFrame], operation: str, parameters: "Optional[Any]" = None, /, **options: Unpack[ExecuteOptions], - ) -> "MockCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": + ) -> "MockCursor[Never, GenericNewPandasDataFrame, Never]": ... @overload def execute_map( self, - row_type: Type[GenericNewPandasDataFlame], + row_type: Type[GenericNewPyArrowTable], operation: str, parameters: "Optional[Any]" = None, /, **options: Unpack[ExecuteOptions], - ) -> "MockCursor[Never, GenericNewPandasDataFlame, Never]": + ) -> "MockCursor[Never, Never, GenericNewPyArrowTable]": ... @overload def execute_map( self, - row_type: Type[GenericNewPyArrowTable], + row_type: Type[GenericNewPanderaDataFrameModel], operation: str, parameters: "Optional[Any]" = None, /, **options: Unpack[ExecuteOptions], - ) -> "MockCursor[Never, Never, GenericNewPyArrowTable]": + ) -> "MockCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": ... @override @@ -105,9 +107,9 @@ 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, @@ -135,34 +137,34 @@ def executemany_map( @overload def executemany_map( self, - row_type: Type[GenericNewPanderaDataFrameModel], + row_type: Type[GenericNewPandasDataFrame], operation: str, seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> "MockCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": + ) -> "MockCursor[Never, GenericNewPandasDataFrame, Never]": ... @overload def executemany_map( self, - row_type: Type[GenericNewPandasDataFlame], + row_type: Type[GenericNewPyArrowTable], operation: str, seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> "MockCursor[Never, GenericNewPandasDataFlame, Never]": + ) -> "MockCursor[Never, Never, GenericNewPyArrowTable]": ... @overload def executemany_map( self, - row_type: Type[GenericNewPyArrowTable], + row_type: Type[GenericNewPanderaDataFrameModel], operation: str, seq_of_parameters: Sequence[Any], /, **options: Unpack[ExecuteOptions], - ) -> "MockCursor[Never, Never, GenericNewPyArrowTable]": + ) -> "MockCursor[Never, PanderaDataFrame[GenericNewPanderaDataFrameModel], Never]": ... @override @@ -170,9 +172,9 @@ 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], @@ -206,7 +208,7 @@ def fetch_arrow_all(self) -> GenericPyArrowTable: def fetch_arrow_batches(self) -> Iterator[GenericPyArrowTable]: return iter([self.fetch_arrow_all()]) - def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFlame: + def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFrame: df = self.fetchone() if ( @@ -216,9 +218,9 @@ def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFlame: ): self._row_type.validate(df) # type: ignore - return cast(GenericPandasDataFlame, df) + return cast(GenericPandasDataFrame, df) - def fetch_pandas_batches(self, **kwargs) -> Iterator[GenericPandasDataFlame]: + def fetch_pandas_batches(self, **kwargs) -> Iterator[GenericPandasDataFrame]: """Fetches a single Arrow Table.""" return iter([self.fetch_pandas_all(**kwargs)]) diff --git a/turu-snowflake/src/turu/snowflake/record/__init__.py b/turu-snowflake/src/turu/snowflake/record/__init__.py index 6f6e69d..19bf559 100644 --- a/turu-snowflake/src/turu/snowflake/record/__init__.py +++ b/turu-snowflake/src/turu/snowflake/record/__init__.py @@ -1,10 +1,10 @@ from turu.core.record.csv_recorder import add_record_map -from turu.snowflake.features import USE_PANDAS, PandasDataFlame +from turu.snowflake.features import USE_PANDAS, PandasDataFrame if USE_PANDAS: add_record_map( "pandas", - lambda row: USE_PANDAS and isinstance(row, PandasDataFlame), + lambda row: USE_PANDAS and isinstance(row, PandasDataFrame), lambda row: row.keys(), lambda row: row.values, lambda row: row, diff --git a/turu-snowflake/src/turu/snowflake/record/async_record_cursor.py b/turu-snowflake/src/turu/snowflake/record/async_record_cursor.py index 8bf440b..8eb029d 100644 --- a/turu-snowflake/src/turu/snowflake/record/async_record_cursor.py +++ b/turu-snowflake/src/turu/snowflake/record/async_record_cursor.py @@ -1,10 +1,10 @@ import turu.core.record -from turu.snowflake.features import PandasDataFlame +from turu.snowflake.features import PandasDataFrame class AsyncRecordCursor(turu.core.record.AsyncRecordCursor): - async def fetch_pandas_all(self, **kwargs) -> "PandasDataFlame": - df: PandasDataFlame = self._raw_cursor.fetch_pandas_all(**kwargs) # type: ignore + async def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame": + df: PandasDataFrame = self._raw_cursor.fetch_pandas_all(**kwargs) # type: ignore if isinstance(self._recorder, turu.core.record.CsvRecorder): df.to_csv(self._recorder.file, index=False) diff --git a/turu-snowflake/src/turu/snowflake/record/record_cursor.py b/turu-snowflake/src/turu/snowflake/record/record_cursor.py index 534e1a4..9b27e0e 100644 --- a/turu-snowflake/src/turu/snowflake/record/record_cursor.py +++ b/turu-snowflake/src/turu/snowflake/record/record_cursor.py @@ -1,10 +1,10 @@ import turu.core.record -from turu.snowflake.features import PandasDataFlame +from turu.snowflake.features import PandasDataFrame class RecordCursor(turu.core.record.RecordCursor): - def fetch_pandas_all(self, **kwargs) -> "PandasDataFlame": - df: PandasDataFlame = self._raw_cursor.fetch_pandas_all(**kwargs) # type: ignore + def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame": + df: PandasDataFrame = self._raw_cursor.fetch_pandas_all(**kwargs) # type: ignore if isinstance(self._recorder, turu.core.record.CsvRecorder): df.to_csv(self._recorder.file, index=False) diff --git a/turu-snowflake/tests/turu/test_snowflake.py b/turu-snowflake/tests/turu/test_snowflake.py index 26c3d81..5543f6a 100644 --- a/turu-snowflake/tests/turu/test_snowflake.py +++ b/turu-snowflake/tests/turu/test_snowflake.py @@ -2,12 +2,17 @@ import tempfile from pathlib import Path from textwrap import dedent -from typing import Annotated, NamedTuple +from typing import Annotated, NamedTuple, cast import pytest import turu.snowflake from turu.core.record import record_to_csv -from turu.snowflake.features import USE_PANDAS, USE_PANDERA, USE_PYARROW +from turu.snowflake.features import ( + USE_PANDAS, + USE_PANDERA, + USE_PYARROW, +) +from typing_extensions import Never def test_version(): @@ -30,6 +35,53 @@ def test_execute(self, connection: turu.snowflake.Connection): def test_execute_fetchone(self, connection: turu.snowflake.Connection): assert connection.execute("select 1").fetchone() == (1,) + def test_execute_map_named_tuple_type(self, connection: turu.snowflake.Connection): + class Row(NamedTuple): + pass + + _cursor: turu.snowflake.Cursor[Row, Never, Never] = connection.execute_map( + Row, "select 1" + ) + + def test_execute_map_dataclass_type(self, connection: turu.snowflake.Connection): + from dataclasses import dataclass + + @dataclass + class Row(NamedTuple): + pass + + _cursor: turu.snowflake.Cursor[Row, Never, Never] = connection.execute_map( + Row, "select 1" + ) + + @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed") + def test_execute_map_pandas_type(self, connection: turu.snowflake.Connection): + import pandas as pd + + _cursor: turu.snowflake.Cursor[ + Never, pd.DataFrame, Never + ] = connection.execute_map(pd.DataFrame, "select 1") + + @pytest.mark.skipif(not USE_PYARROW, reason="pyarrow is not installed") + def test_execute_pyarrow_type(self, connection: turu.snowflake.Connection): + from turu.snowflake.features import PyArrowTable + + _cursor: turu.snowflake.Cursor[ + Never, Never, PyArrowTable + ] = connection.execute_map(PyArrowTable, "select 1") + + @pytest.mark.skipif(not USE_PANDERA, reason="pandera is not installed") + def test_execute_map_pandera_type(self, connection: turu.snowflake.Connection): + import pandera as pa # type: ignore[import] + from turu.snowflake.features import PanderaDataFrame + + class RowModel(pa.DataFrameModel): + ID: pa.Int8 + + _cursor: turu.snowflake.Cursor[ + Never, PanderaDataFrame[RowModel], Never + ] = connection.execute_map(RowModel, "select 1 as ID") + def test_execute_map_fetchone(self, connection: turu.snowflake.Connection): cursor = connection.execute_map(Row, "select 1") @@ -169,7 +221,7 @@ def test_fetch_arrow_batches(self, connection: turu.snowflake.Connection): with connection.execute("select 1 as ID union all select 2 as ID") as cursor: for row in cursor.fetch_arrow_batches(): assert_frame_equal( - row.to_pandas(), + cast(DataFrame, row.to_pandas()), DataFrame({"ID": [1, 2]}, dtype="int8"), ) diff --git a/turu-snowflake/tests/turu/test_snowflake_async.py b/turu-snowflake/tests/turu/test_snowflake_async.py index 97782e6..9078696 100644 --- a/turu-snowflake/tests/turu/test_snowflake_async.py +++ b/turu-snowflake/tests/turu/test_snowflake_async.py @@ -2,12 +2,13 @@ import tempfile from pathlib import Path from textwrap import dedent -from typing import Annotated, NamedTuple +from typing import Annotated, NamedTuple, cast import pytest import turu.snowflake from turu.core.record import record_to_csv from turu.snowflake.features import USE_PANDAS, USE_PANDERA, USE_PYARROW +from typing_extensions import Never class TestTuruSnowflake: @@ -30,6 +31,68 @@ async def test_execute(self, async_connection: turu.snowflake.AsyncConnection): cursor = await async_connection.execute("select 1") assert await cursor.fetchall() == [(1,)] + @pytest.mark.asyncio + async def test_execute_map_named_tuple_type( + self, async_connection: turu.snowflake.AsyncConnection + ): + class Row(NamedTuple): + pass + + _cursor: turu.snowflake.AsyncCursor[ + Row, Never, Never + ] = await async_connection.execute_map(Row, "select 1") + + @pytest.mark.asyncio + async def test_execute_map_dataclass_type( + self, async_connection: turu.snowflake.AsyncConnection + ): + from dataclasses import dataclass + + @dataclass + class Row(NamedTuple): + pass + + _cursor: turu.snowflake.AsyncCursor[ + Row, Never, Never + ] = await async_connection.execute_map(Row, "select 1") + + @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed") + @pytest.mark.asyncio + async def test_execute_map_pandas_type( + self, async_connection: turu.snowflake.AsyncConnection + ): + import pandas as pd + + _cursor: turu.snowflake.AsyncCursor[ + Never, pd.DataFrame, Never + ] = await async_connection.execute_map(pd.DataFrame, "select 1") + + @pytest.mark.skipif(not USE_PYARROW, reason="pyarrow is not installed") + @pytest.mark.asyncio + async def test_execute_pyarrow_type( + self, async_connection: turu.snowflake.AsyncConnection + ): + import pyarrow as pa # type: ignore[import] + + _cursor: turu.snowflake.AsyncCursor[ + Never, Never, pa.Table + ] = await async_connection.execute_map(pa.Table, "select 1") + + @pytest.mark.skipif(not USE_PANDERA, reason="pandera is not installed") + @pytest.mark.asyncio + async def test_execute_map_pandera_type( + self, async_connection: turu.snowflake.AsyncConnection + ): + import pandera as pa # type: ignore[import] + from turu.snowflake.features import PanderaDataFrame + + class RowModel(pa.DataFrameModel): + ID: pa.Int8 + + _cursor: turu.snowflake.AsyncCursor[ + Never, PanderaDataFrame[RowModel], Never + ] = await async_connection.execute_map(RowModel, "select 1 as ID") + @pytest.mark.asyncio async def test_execute_fetchone( self, async_connection: turu.snowflake.AsyncConnection @@ -264,7 +327,8 @@ async def test_fetch_arrow_batches( ) as cursor: for row in await cursor.fetch_arrow_batches(): assert_frame_equal( - row.to_pandas(), DataFrame({"ID": [1, 2]}, dtype="int8") + cast(DataFrame, row.to_pandas()), + DataFrame({"ID": [1, 2]}, dtype="int8"), ) @pytest.mark.skipif( diff --git a/turu-snowflake/tests/turu/test_snowflake_mock.py b/turu-snowflake/tests/turu/test_snowflake_mock.py index 360cd99..8669039 100644 --- a/turu-snowflake/tests/turu/test_snowflake_mock.py +++ b/turu-snowflake/tests/turu/test_snowflake_mock.py @@ -9,9 +9,10 @@ USE_PANDAS, USE_PANDERA, USE_PYARROW, - PandasDataFlame, + PandasDataFrame, PyArrowTable, ) +from typing_extensions import Never class Row(NamedTuple): @@ -25,6 +26,73 @@ def test_execute(self, mock_connection: turu.snowflake.MockConnection): assert cursor.fetchone() == (1,) assert cursor.fetchone() is None + def test_execute_map_named_tuple_type( + self, mock_connection: turu.snowflake.MockConnection + ): + class Row(NamedTuple): + id: int + + _cursor: turu.snowflake.Cursor[ + Row, Never, Never + ] = mock_connection.inject_response(Row, Row(1)).execute_map(Row, "select 1") + + def test_execute_map_dataclass_type( + self, mock_connection: turu.snowflake.MockConnection + ): + from dataclasses import dataclass + + @dataclass + class Row: + id: int + + _cursor: turu.snowflake.Cursor[ + Row, Never, Never + ] = mock_connection.inject_response(Row, Row(id=1)).execute_map(Row, "select 1") + + @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed") + def test_execute_map_pandas_type( + self, mock_connection: turu.snowflake.MockConnection + ): + import pandas as pd + + _cursor: turu.snowflake.Cursor[ + Never, pd.DataFrame, Never + ] = mock_connection.inject_response( + pd.DataFrame, pd.DataFrame({"id": [1]}) + ).execute_map(pd.DataFrame, "select 1") + + @pytest.mark.skipif(not USE_PYARROW, reason="pyarrow is not installed") + def test_execute_pyarrow_type(self, mock_connection: turu.snowflake.MockConnection): + import pyarrow as pa # type: ignore[import] + + expected: pa.Table = pa.table( + data=[pa.array([1], type=pa.int64())], + schema=pa.schema([pa.field("ID", pa.int64())]), + ) # type: ignore + _cursor: turu.snowflake.Cursor[ + Never, Never, pa.Table + ] = mock_connection.inject_response( + pa.Table, + expected, + ).execute_map(pa.Table, "select 1") + + @pytest.mark.skipif(not USE_PANDERA, reason="pandera is not installed") + def test_execute_map_pandera_type( + self, mock_connection: turu.snowflake.MockConnection + ): + import pandas as pd # type: ignore[import] + import pandera as pa # type: ignore[import] + from turu.snowflake.features import PanderaDataFrame + + class RowModel(pa.DataFrameModel): + ID: pa.Int8 + + _cursor: turu.snowflake.Cursor[ + Never, PanderaDataFrame[RowModel], Never + ] = mock_connection.inject_response( + RowModel, pd.DataFrame({"id": [1]}) + ).execute_map(RowModel, "select 1 as ID") + def test_execute_map_fetchone(self, mock_connection: turu.snowflake.MockConnection): expected = [Row(1), Row(2)] @@ -208,10 +276,10 @@ def test_fetch_pandas_all(self, mock_connection: turu.snowflake.MockConnection): expected = pd.DataFrame({"ID": [1, 2]}) - mock_connection.inject_response(PandasDataFlame, expected) + mock_connection.inject_response(PandasDataFrame, expected) with mock_connection.execute_map( - PandasDataFlame, "select 1 as ID union all select 2 as ID" + PandasDataFrame, "select 1 as ID union all select 2 as ID" ) as cursor: assert cursor.fetch_pandas_all().equals(expected) @@ -221,8 +289,8 @@ def test_fetch_pandas_batches(self, mock_connection: turu.snowflake.MockConnecti expected = pd.DataFrame({"ID": [1, 2]}) - with mock_connection.inject_response(PandasDataFlame, expected).execute_map( - PandasDataFlame, "select 1 as ID union all select 2 as ID" + with mock_connection.inject_response(PandasDataFrame, expected).execute_map( + PandasDataFrame, "select 1 as ID union all select 2 as ID" ) as cursor: assert list(cursor.fetch_pandas_batches()) == [expected] diff --git a/turu-snowflake/tests/turu/test_snowflake_mock_async.py b/turu-snowflake/tests/turu/test_snowflake_mock_async.py index d676314..acf1254 100644 --- a/turu-snowflake/tests/turu/test_snowflake_mock_async.py +++ b/turu-snowflake/tests/turu/test_snowflake_mock_async.py @@ -6,6 +6,7 @@ import pytest import turu.snowflake from turu.snowflake.features import USE_PANDAS, USE_PANDERA, USE_PYARROW, PyArrowTable +from typing_extensions import Never class Row(NamedTuple): @@ -22,6 +23,84 @@ async def test_execute( assert await cursor.fetchone() == (1,) assert await cursor.fetchone() is None + @pytest.mark.asyncio + async def test_execute_map_named_tuple_type( + self, mock_async_connection: turu.snowflake.MockAsyncConnection + ): + class Row(NamedTuple): + id: int + + _cursor: turu.snowflake.AsyncCursor[ + Row, Never, Never + ] = await mock_async_connection.inject_response(Row, Row(1)).execute_map( + Row, "select 1" + ) + + @pytest.mark.asyncio + async def test_execute_map_dataclass_type( + self, mock_async_connection: turu.snowflake.MockAsyncConnection + ): + from dataclasses import dataclass + + @dataclass + class Row: + id: int + + _cursor: turu.snowflake.AsyncCursor[ + Row, Never, Never + ] = await mock_async_connection.inject_response(Row, Row(id=1)).execute_map( + Row, "select 1" + ) + + @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed") + @pytest.mark.asyncio + async def test_execute_map_pandas_type( + self, mock_async_connection: turu.snowflake.MockAsyncConnection + ): + import pandas as pd + + _cursor: turu.snowflake.AsyncCursor[ + Never, pd.DataFrame, Never + ] = await mock_async_connection.inject_response( + pd.DataFrame, pd.DataFrame({"id": [1]}) + ).execute_map(pd.DataFrame, "select 1") + + @pytest.mark.skipif(not USE_PYARROW, reason="pyarrow is not installed") + @pytest.mark.asyncio + async def test_execute_pyarrow_type( + self, mock_async_connection: turu.snowflake.MockAsyncConnection + ): + import pyarrow as pa # type: ignore[import] + + expected: pa.Table = pa.table( + data=[pa.array([1], type=pa.int64())], + schema=pa.schema([pa.field("ID", pa.int64())]), + ) # type: ignore + _cursor: turu.snowflake.AsyncCursor[ + Never, Never, pa.Table + ] = await mock_async_connection.inject_response( + pa.Table, + expected, + ).execute_map(pa.Table, "select 1") + + @pytest.mark.skipif(not USE_PANDERA, reason="pandera is not installed") + @pytest.mark.asyncio + async def test_execute_map_pandera_type( + self, mock_async_connection: turu.snowflake.MockAsyncConnection + ): + import pandas as pd # type: ignore[import] + import pandera as pa # type: ignore[import] + from turu.snowflake.features import PanderaDataFrame + + class RowModel(pa.DataFrameModel): + ID: pa.Int8 + + _cursor: turu.snowflake.AsyncCursor[ + Never, PanderaDataFrame[RowModel], Never + ] = await mock_async_connection.inject_response( + RowModel, pd.DataFrame({"id": [1]}) + ).execute_map(RowModel, "select 1 as ID") + @pytest.mark.asyncio async def test_execute_map_fetchone( self, mock_async_connection: turu.snowflake.MockAsyncConnection