Fix data type bit width

kylebarron · Sep 22, 2024 · 2425d7d · 2425d7d
1 parent daecc97
commit 2425d7d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/arro3-core/python/arro3/core/_core.pyi b/arro3-core/python/arro3/core/_core.pyi
@@ -355,7 +355,11 @@ class DataType:
     def from_arrow_pycapsule(cls, capsule) -> DataType:
         """Construct this object from a bare Arrow PyCapsule"""
     @property
-    def bit_width(self) -> int | None: ...
+    def bit_width(self) -> Literal[8, 16, 32, 64] | None:
+        """Returns the bit width of this type if it is a primitive type
+
+        Returns `None` if not a primitive type
+        """
     def equals(
         self, other: ArrowSchemaExportable, *, check_metadata: bool = False
     ) -> bool:

diff --git a/pyo3-arrow/src/datatypes.rs b/pyo3-arrow/src/datatypes.rs
@@ -113,6 +113,7 @@ impl Display for PyDataType {
     }
 }
 
+#[allow(non_snake_case)]
 #[pymethods]
 impl PyDataType {
     fn __arrow_c_schema__<'py>(&'py self, py: Python<'py>) -> PyArrowResult<Bound<'py, PyCapsule>> {
@@ -140,7 +141,7 @@ impl PyDataType {
 
     #[getter]
     fn bit_width(&self) -> Option<usize> {
-        self.0.primitive_width()
+        self.0.primitive_width().map(|width| width * 8)
     }
 
     #[pyo3(signature=(other, *, check_metadata=false))]