diff --git a/protocol/dataframe_protocol.py b/protocol/dataframe_protocol.py index 4d9e6c19..209072be 100644 --- a/protocol/dataframe_protocol.py +++ b/protocol/dataframe_protocol.py @@ -58,6 +58,9 @@ class DtypeKind(enum.IntEnum): CATEGORICAL = 23 +Dtype = Tuple[DtypeKind, int, str, str] # see Column.dtype + + class ColumnNullType(enum.IntEnum): """ Integer enum for null type representation. @@ -86,18 +89,18 @@ class ColumnNullType(enum.IntEnum): class ColumnBuffers(TypedDict): # first element is a buffer containing the column data; # second element is the data buffer's associated dtype - data: Tuple["Buffer", Any] + data: Tuple["Buffer", Dtype] # first element is a buffer containing mask values indicating missing data; # second element is the mask value buffer's associated dtype. # None if the null representation is not a bit or byte mask - validity: Optional[Tuple["Buffer", Any]] + validity: Optional[Tuple["Buffer", Dtype]] # first element is a buffer containing the offset values for # variable-size binary data (e.g., variable-length strings); # second element is the offsets buffer's associated dtype. # None if the data buffer does not have an associated offsets buffer - offsets: Optional[Tuple["Buffer", Any]] + offsets: Optional[Tuple["Buffer", Dtype]] class CategoricalDescription(TypedDict): @@ -237,7 +240,7 @@ def offset(self) -> int: @property @abstractmethod - def dtype(self) -> Tuple[DtypeKind, int, str, str]: + def dtype(self) -> Dtype: """ Dtype description as a tuple ``(kind, bit-width, format string, endianness)``.