🏷️ typing

Author: MarcoGorelli

pandas/core/interchange/column.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import (
+    TYPE_CHECKING,
+    Any,
+)
 
 import numpy as np
 
@@ -33,6 +36,9 @@
     dtype_to_arrow_c_fmt,
 )
 
+if TYPE_CHECKING:
+    from pandas.core.interchange.dataframe_protocol import Buffer
+
 _NP_KINDS = {
     "i": DtypeKind.INT,
     "u": DtypeKind.UINT,
@@ -296,7 +302,7 @@ def get_buffers(self) -> ColumnBuffers:
 
     def _get_data_buffer(
         self,
-    ) -> tuple[PandasBuffer, Any]:  # Any is for self.dtype tuple
+    ) -> tuple[Buffer, tuple[DtypeKind, int, str, str]]:
         """
         Return the buffer containing the data and the buffer's associated dtype.
         """
@@ -307,7 +313,7 @@ def _get_data_buffer(
                 np_arr = self._col.dt.tz_convert(None).to_numpy()
             else:
                 np_arr = self._col.to_numpy()
-            buffer = PandasBuffer(np_arr, allow_copy=self._allow_copy)
+            buffer: Buffer = PandasBuffer(np_arr, allow_copy=self._allow_copy)
             dtype = (
                 DtypeKind.INT,
                 64,
@@ -324,7 +330,9 @@ def _get_data_buffer(
             arr = self._col.array
             if isinstance(self._col.dtype, ArrowDtype):
                 buffer = PandasBufferPyarrow(
-                    arr._pa_array, is_validity=False, allow_copy=self._allow_copy
+                    arr._pa_array,  # type: ignore[attr-defined]
+                    is_validity=False,
+                    allow_copy=self._allow_copy,
                 )
                 if self.dtype[0] == DtypeKind.BOOL:
                     dtype = (
@@ -371,7 +379,7 @@ def _get_data_buffer(
 
         return buffer, dtype
 
-    def _get_validity_buffer(self) -> tuple[PandasBuffer, Any]:
+    def _get_validity_buffer(self) -> tuple[Buffer, Any] | None:
         """
         Return the buffer containing the mask values indicating missing data and
         the buffer's associated dtype.
@@ -382,10 +390,15 @@ def _get_validity_buffer(self) -> tuple[PandasBuffer, Any]:
         if isinstance(self._col.dtype, ArrowDtype):
             arr = self._col.array
             dtype = (DtypeKind.BOOL, 1, ArrowCTypes.BOOL, Endianness.NATIVE)
-            if all(chunk.buffers()[0] is None for chunk in arr._pa_array.chunks):
+            if all(
+                chunk.buffers()[0] is None
+                for chunk in arr._pa_array.chunks  # type: ignore[attr-defined]
+            ):
                 return None
-            buffer = PandasBufferPyarrow(
-                arr._pa_array, is_validity=True, allow_copy=self._allow_copy
+            buffer: Buffer = PandasBufferPyarrow(
+                arr._pa_array,  # type: ignore[attr-defined]
+                is_validity=True,
+                allow_copy=self._allow_copy,
             )
             return buffer, dtype
 

pandas/tests/interchange/test_impl.py

@@ -552,7 +552,9 @@ def test_string_validity_buffer() -> None:
 def test_string_validity_buffer_no_missing() -> None:
     # https://github.com/pandas-dev/pandas/issues/57762
     df = pd.DataFrame({"a": ["x", None]}, dtype="large_string[pyarrow]")
-    result = df.__dataframe__().get_column_by_name("a").get_buffers()["validity"][1]
+    validity = df.__dataframe__().get_column_by_name("a").get_buffers()["validity"]
+    assert validity is not None
+    result = validity[1]
     expected = (DtypeKind.BOOL, 1, ArrowCTypes.BOOL, "=")
     assert result == expected