[python-package] fix mypy error about pyarrow data (#6333)

microsoft · Feb 21, 2024 · b60068c · b60068c
1 parent 008ebaf
commit b60068c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/.ci/test.sh b/.ci/test.sh
@@ -76,6 +76,7 @@ if [[ $TASK == "lint" ]]; then
         cpplint \
         mypy \
         'pre-commit>=3.6.0' \
+        'pyarrow>=14.0' \
         'r-lintr>=3.1'
     source activate $CONDA_ENV
     echo "Linting Python code"

diff --git a/python-package/lightgbm/basic.py b/python-package/lightgbm/basic.py
@@ -364,7 +364,7 @@ def _is_2d_collection(data: Any) -> bool:
     )
 
 
-def _is_pyarrow_array(data: Any) -> bool:
+def _is_pyarrow_array(data: Any) -> "TypeGuard[Union[pa_Array, pa_ChunkedArray]]":
     """Check whether data is a PyArrow array."""
     return isinstance(data, (pa_Array, pa_ChunkedArray))