fix: typehint.

Author: ya7010

turu-snowflake/poetry.lock

@@ -63,6 +63,7 @@ pytest-cov = "^4.1.0"
 pytest-asyncio = "^0.23.2"
 pandas = "^2.1.4"
 pyarrow = "^14.0.2"
+pyarrow-stubs = "^10.0.1.7"
 numpy = "^1.26.3"
 
 [tool.taskipy.tasks]

turu-snowflake/pyproject.toml

@@ -290,12 +290,17 @@ async def fetchall(self) -> List[GenericRowType]:
     async def fetch_arrow_all(self) -> GenericPyArrowTable:
         """Fetches a single Arrow Table."""
 
-        return self._raw_cursor.fetch_arrow_all(force_return_table=True)
+        return cast(
+            GenericPyArrowTable,
+            self._raw_cursor.fetch_arrow_all(force_return_table=True),
+        )
 
     async def fetch_arrow_batches(self) -> Iterator[GenericPyArrowTable]:
         """Fetches Arrow Tables in batches, where 'batch' refers to Snowflake Chunk."""
 
-        return self._raw_cursor.fetch_arrow_batches()
+        return cast(
+            Iterator[GenericPyArrowTable], self._raw_cursor.fetch_arrow_batches()
+        )
 
     async def fetch_pandas_all(self, **kwargs) -> GenericPandasDataFlame:
         """Fetch Pandas dataframes."""

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

@@ -324,7 +324,9 @@ def fetch_arrow_all(self) -> GenericPyArrowTable:
     def fetch_arrow_batches(self) -> "Iterator[GenericPyArrowTable]":
         """Fetches Arrow Tables in batches, where 'batch' refers to Snowflake Chunk."""
 
-        return self._raw_cursor.fetch_arrow_batches()
+        return cast(
+            Iterator[GenericPyArrowTable], self._raw_cursor.fetch_arrow_batches()
+        )
 
     def fetch_pandas_all(self, **kwargs) -> "GenericPandasDataFlame":
         """Fetch a single Pandas dataframe."""

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

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, TypeVar
+from typing import TypeVar
 
 from turu.core.features import _NotSupportFeature
 from typing_extensions import Never, TypeAlias
@@ -20,11 +20,8 @@
     import pyarrow  # type: ignore[import]  # noqa: F401
 
     USE_PYARROW = True
-    if TYPE_CHECKING:
-        PyArrowTable: TypeAlias = _NotSupportFeature  # type: ignore
+    PyArrowTable = pyarrow.Table  # type: ignore
 
-    else:
-        PyArrowTable: TypeAlias = pyarrow.Table
 
 except ImportError:
     USE_PYARROW = False

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

@@ -127,7 +127,7 @@ def inject_response_from_csv(  # type: ignore[override]
 
                 self.inject_response(
                     row_type,
-                    cast(Any, pyarrow.csv.read_csv(filepath, **options)),
+                    pyarrow.csv.read_csv(filepath, **options),  # type: ignore
                 )
 
             else:

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

@@ -153,10 +153,10 @@ def test_cursor_use_role(self, connection: turu.snowflake.Connection):
     def test_fetch_arrow_all(self, connection: turu.snowflake.Connection):
         import pyarrow as pa
 
-        expected = pa.table(
+        expected: pa.Table = pa.table(
             data=[pa.array([1, 2], type=pa.int8())],
             schema=pa.schema([pa.field("ID", pa.int8(), False)]),
-        )
+        )  # type: ignore
 
         with connection.execute("select 1 as ID union all select 2 as ID") as cursor:
             assert cursor.fetch_arrow_all() == expected

turu-snowflake/tests/turu/test_snowflake.py

@@ -173,10 +173,10 @@ def test_cursor_use_role(self, mock_connection: turu.snowflake.MockConnection):
     def test_fetch_arrow_all(self, mock_connection: turu.snowflake.MockConnection):
         import pyarrow as pa
 
-        expected = pa.table(
+        expected: pa.Table = pa.table(
             data=[pa.array([1, 2], type=pa.int8())],
             schema=pa.schema([pa.field("ID", pa.int8(), False)]),
-        )
+        )  # type: ignore
 
         mock_connection.inject_response(PyArrowTable, expected)
 
@@ -192,10 +192,10 @@ def test_fetch_arrow_all(self, mock_connection: turu.snowflake.MockConnection):
     def test_fetch_arrow_batches(self, mock_connection: turu.snowflake.MockConnection):
         import pyarrow as pa
 
-        expected = pa.table(
+        expected: pa.Table = pa.table(
             data=[pa.array([1, 2], type=pa.int8())],
             schema=pa.schema([pa.field("ID", pa.int8(), False)]),
-        )
+        )  # type: ignore
 
         with mock_connection.inject_response(PyArrowTable, expected).execute_map(
             PyArrowTable, "select 1 as ID union all select 2 as ID"
@@ -232,10 +232,10 @@ def test_inject_pyarrow_response_from_csv(
     ):
         import pyarrow as pa
 
-        expected = pa.table(
+        expected: pa.Table = pa.table(
             data=[pa.array([1, 2], type=pa.int64())],
             schema=pa.schema([pa.field("ID", pa.int64())]),
-        )
+        )  # type: ignore
 
         with tempfile.NamedTemporaryFile() as file:
             Path(file.name).write_text(

turu-snowflake/tests/turu/test_snowflake_mock.py

@@ -229,10 +229,10 @@ async def test_fetch_arrow_all(
     ):
         import pyarrow as pa
 
-        expected = pa.table(
+        expected: pa.Table = pa.table(
             data=[pa.array([1, 2], type=pa.int8())],
             schema=pa.schema([pa.field("ID", pa.int8(), False)]),
-        )
+        )  # type: ignore
 
         mock_async_connection.inject_response(PyArrowTable, expected)
 
@@ -251,10 +251,10 @@ async def test_fetch_arrow_batches(
     ):
         import pyarrow as pa
 
-        expected = pa.table(
+        expected: pa.Table = pa.table(
             data=[pa.array([1, 2], type=pa.int8())],
             schema=pa.schema([pa.field("ID", pa.int8(), False)]),
-        )
+        )  # type: ignore
 
         async with await mock_async_connection.inject_response(
             PyArrowTable, expected
@@ -344,10 +344,10 @@ async def test_inject_pyarrow_response_from_csv(
     ):
         import pyarrow as pa
 
-        expected = pa.table(
+        expected: pa.Table = pa.table(
             data=[pa.array([1, 2], type=pa.int64())],
             schema=pa.schema([pa.field("ID", pa.int64())]),
-        )
+        )  # type: ignore
 
         with tempfile.NamedTemporaryFile() as file:
             Path(file.name).write_text(