From 137ed2998f1854f2f1f5fa2e338abb98872a81f6 Mon Sep 17 00:00:00 2001 From: Kyle Barron Date: Wed, 10 Jul 2024 17:41:04 -0400 Subject: [PATCH] implement basic `__repr__` on classes (#48) --- pyo3-arrow/src/array.rs | 14 ++++++++++++++ pyo3-arrow/src/chunked.rs | 14 ++++++++++++++ pyo3-arrow/src/field.rs | 19 +++++++++++++++++++ pyo3-arrow/src/record_batch.rs | 14 ++++++++++++++ pyo3-arrow/src/record_batch_reader.rs | 18 ++++++++++++++++++ pyo3-arrow/src/schema.rs | 24 ++++++++++++++++++++++++ pyo3-arrow/src/table.rs | 14 ++++++++++++++ 7 files changed, 117 insertions(+) diff --git a/pyo3-arrow/src/array.rs b/pyo3-arrow/src/array.rs index 51031dd..99f370e 100644 --- a/pyo3-arrow/src/array.rs +++ b/pyo3-arrow/src/array.rs @@ -1,4 +1,5 @@ use std::ffi::CString; +use std::fmt::Display; use std::sync::Arc; use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema}; @@ -95,6 +96,15 @@ impl AsRef for PyArray { } } +impl Display for PyArray { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "arro3.Array<")?; + self.array.data_type().fmt(f)?; + writeln!(f, ">")?; + Ok(()) + } +} + #[pymethods] impl PyArray { /// An implementation of the Array interface, for interoperability with numpy and other @@ -138,6 +148,10 @@ impl PyArray { self.array.len() } + pub fn __repr__(&self) -> String { + self.to_string() + } + /// Construct this object from an existing Arrow object. /// /// It can be called on anything that exports the Arrow data interface diff --git a/pyo3-arrow/src/chunked.rs b/pyo3-arrow/src/chunked.rs index 36661d0..330fd3d 100644 --- a/pyo3-arrow/src/chunked.rs +++ b/pyo3-arrow/src/chunked.rs @@ -1,4 +1,5 @@ use std::ffi::CString; +use std::fmt::Display; use std::sync::Arc; use arrow_array::{make_array, Array, ArrayRef}; @@ -114,6 +115,15 @@ impl AsRef<[ArrayRef]> for PyChunkedArray { } } +impl Display for PyChunkedArray { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "arro3.ChunkedArray<")?; + self.field.data_type().fmt(f)?; + writeln!(f, ">")?; + Ok(()) + } +} + #[pymethods] impl PyChunkedArray { /// An implementation of the Array interface, for interoperability with numpy and other @@ -158,6 +168,10 @@ impl PyChunkedArray { self.chunks.iter().fold(0, |acc, x| acc + x.len()) } + pub fn __repr__(&self) -> String { + self.to_string() + } + /// Copy this array to a `numpy` NDArray pub fn to_numpy(&self, py: Python) -> PyResult { self.__array__(py) diff --git a/pyo3-arrow/src/field.rs b/pyo3-arrow/src/field.rs index 7d56d91..cec35a5 100644 --- a/pyo3-arrow/src/field.rs +++ b/pyo3-arrow/src/field.rs @@ -1,4 +1,5 @@ use std::ffi::CString; +use std::fmt::Display; use std::sync::Arc; use arrow::ffi::FFI_ArrowSchema; @@ -68,6 +69,20 @@ impl AsRef for PyField { } } +impl Display for PyField { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "arro3.Field<")?; + f.write_str(self.0.name().as_str())?; + write!(f, ": ")?; + self.0.data_type().fmt(f)?; + if !self.0.is_nullable() { + write!(f, " not null")?; + } + writeln!(f, ">")?; + Ok(()) + } +} + #[pymethods] impl PyField { /// An implementation of the [Arrow PyCapsule @@ -91,6 +106,10 @@ impl PyField { self.0 == other.0 } + pub fn __repr__(&self) -> String { + self.to_string() + } + /// Construct this from an existing Arrow object. /// /// It can be called on anything that exports the Arrow schema interface diff --git a/pyo3-arrow/src/record_batch.rs b/pyo3-arrow/src/record_batch.rs index d03e202..e0cd727 100644 --- a/pyo3-arrow/src/record_batch.rs +++ b/pyo3-arrow/src/record_batch.rs @@ -1,4 +1,5 @@ use std::ffi::CString; +use std::fmt::Display; use std::sync::Arc; use arrow::array::AsArray; @@ -13,6 +14,7 @@ use pyo3::types::{PyCapsule, PyTuple, PyType}; use crate::error::PyArrowResult; use crate::ffi::from_python::utils::import_array_pycapsules; use crate::ffi::to_python::nanoarrow::to_nanoarrow_array; +use crate::schema::display_schema; /// A Python-facing Arrow record batch. /// @@ -77,6 +79,14 @@ impl AsRef for PyRecordBatch { } } +impl Display for PyRecordBatch { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "arro3.RecordBatch")?; + writeln!(f, "-----------------")?; + display_schema(&self.0.schema(), f) + } +} + #[pymethods] impl PyRecordBatch { /// An implementation of the [Arrow PyCapsule @@ -109,6 +119,10 @@ impl PyRecordBatch { self.0 == other.0 } + pub fn __repr__(&self) -> String { + self.to_string() + } + /// Construct this from an existing Arrow RecordBatch. /// /// It can be called on anything that exports the Arrow data interface diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs index 34c91dc..055f6af 100644 --- a/pyo3-arrow/src/record_batch_reader.rs +++ b/pyo3-arrow/src/record_batch_reader.rs @@ -1,4 +1,5 @@ use std::ffi::CString; +use std::fmt::Display; use arrow::ffi_stream::FFI_ArrowArrayStream; use arrow_array::RecordBatchReader; @@ -11,6 +12,7 @@ use pyo3::types::{PyCapsule, PyTuple, PyType}; use crate::error::PyArrowResult; use crate::ffi::from_python::utils::import_stream_pycapsule; use crate::ffi::to_python::nanoarrow::to_nanoarrow_array_stream; +use crate::schema::display_schema; use crate::{PySchema, PyTable}; /// A Python-facing Arrow record batch reader. @@ -97,6 +99,18 @@ impl From> for PyRecordBatchReader { } } +impl Display for PyRecordBatchReader { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "arro3.RecordBatchReader")?; + writeln!(f, "-----------------------")?; + if let Ok(schema) = self.schema_ref() { + display_schema(&schema, f) + } else { + writeln!(f, "Closed stream") + } + } +} + #[pymethods] impl PyRecordBatchReader { /// An implementation of the [Arrow PyCapsule @@ -122,6 +136,10 @@ impl PyRecordBatchReader { PyCapsule::new_bound(py, ffi_stream, Some(stream_capsule_name)) } + pub fn __repr__(&self) -> String { + self.to_string() + } + /// Returns `true` if this reader has already been consumed. pub fn closed(&self) -> bool { self.0.is_none() diff --git a/pyo3-arrow/src/schema.rs b/pyo3-arrow/src/schema.rs index 305afe9..4f749e5 100644 --- a/pyo3-arrow/src/schema.rs +++ b/pyo3-arrow/src/schema.rs @@ -1,4 +1,5 @@ use std::ffi::CString; +use std::fmt::Display; use std::sync::Arc; use arrow::ffi::FFI_ArrowSchema; @@ -68,6 +69,25 @@ impl AsRef for PySchema { } } +impl Display for PySchema { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "arro3.Schema")?; + writeln!(f, "------------")?; + display_schema(&self.0, f) + } +} + +pub(crate) fn display_schema(schema: &Schema, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + schema.fields().iter().try_for_each(|field| { + f.write_str(field.name().as_str())?; + write!(f, ": ")?; + field.data_type().fmt(f)?; + writeln!(f)?; + Ok::<_, std::fmt::Error>(()) + })?; + Ok(()) +} + #[pymethods] impl PySchema { /// An implementation of the [Arrow PyCapsule @@ -84,6 +104,10 @@ impl PySchema { Ok(schema_capsule) } + pub fn __repr__(&self) -> String { + self.to_string() + } + /// Construct this object from an existing Arrow object /// /// It can be called on anything that exports the Arrow data interface diff --git a/pyo3-arrow/src/table.rs b/pyo3-arrow/src/table.rs index 176cb1b..c342ae0 100644 --- a/pyo3-arrow/src/table.rs +++ b/pyo3-arrow/src/table.rs @@ -1,4 +1,5 @@ use std::ffi::CString; +use std::fmt::Display; use arrow::ffi_stream::ArrowArrayStreamReader as ArrowRecordBatchStreamReader; use arrow::ffi_stream::FFI_ArrowArrayStream; @@ -12,6 +13,7 @@ use pyo3::types::{PyCapsule, PyTuple, PyType}; use crate::ffi::from_python::utils::import_stream_pycapsule; use crate::ffi::to_python::nanoarrow::to_nanoarrow_array_stream; +use crate::schema::display_schema; use crate::PySchema; /// A Python-facing Arrow table. @@ -64,6 +66,14 @@ impl PyTable { } } +impl Display for PyTable { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + writeln!(f, "arro3.Table")?; + writeln!(f, "-----------")?; + display_schema(&self.schema, f) + } +} + #[pymethods] impl PyTable { /// An implementation of the [Arrow PyCapsule @@ -99,6 +109,10 @@ impl PyTable { self.batches.iter().fold(0, |acc, x| acc + x.num_rows()) } + pub fn __repr__(&self) -> String { + self.to_string() + } + /// Construct this object from an existing Arrow object. /// /// It can be called on anything that exports the Arrow stream interface