From 46e70094af11f92d7a6f64069118712171e8bd55 Mon Sep 17 00:00:00 2001 From: Walnut <39544927+Walnut356@users.noreply.github.com> Date: Wed, 18 Oct 2023 01:47:29 -0500 Subject: [PATCH] fix(rust, python): Edge cases for list count formatting (#11780) --- crates/polars-core/src/fmt.rs | 168 +++++++++++++++++++++++++++++----- py-polars/polars/config.py | 4 +- 2 files changed, 146 insertions(+), 26 deletions(-) diff --git a/crates/polars-core/src/fmt.rs b/crates/polars-core/src/fmt.rs index 24758ca320aa..6dec1221081e 100644 --- a/crates/polars-core/src/fmt.rs +++ b/crates/polars-core/src/fmt.rs @@ -964,36 +964,46 @@ fn fmt_struct(f: &mut Formatter<'_>, vals: &[AnyValue]) -> fmt::Result { impl Series { pub fn fmt_list(&self) -> String { - match self.len() { - 0 => "[]".to_string(), - 1 => format!("[{}]", self.get(0).unwrap()), - 2 => format!("[{}, {}]", self.get(0).unwrap(), self.get(1).unwrap()), - 3 => format!( - "[{}, {}, {}]", - self.get(0).unwrap(), - self.get(1).unwrap(), - self.get(2).unwrap() - ), - _ => { - let max_items = std::env::var(FMT_TABLE_CELL_LIST_LEN) - .as_deref() - .unwrap_or("") - .parse() - .map_or(3, |n: i64| if n < 0 { self.len() } else { n as usize }); + if self.is_empty() { + return "[]".to_owned(); + } + + let max_items = std::env::var(FMT_TABLE_CELL_LIST_LEN) + .as_deref() + .unwrap_or("") + .parse() + .map_or(3, |n: i64| if n < 0 { self.len() } else { n as usize }); + match max_items { + 0 => "[…]".to_owned(), + _ if max_items >= self.len() => { let mut result = "[".to_owned(); - for (i, item) in self.iter().enumerate() { + for i in 0..self.len() { + let item = self.get(i).unwrap(); write!(result, "{item}").unwrap(); + // this will always leave a trailing ", " after the last item + // but for long lists, this is faster than checking against the length each time + result.push_str(", "); + } + // remove trailing ", " and replace with closing brace + result.pop(); + result.pop(); + result.push(']'); - if i != self.len() - 1 { - result.push_str(", "); - } + result + }, + _ => { + let mut result = "[".to_owned(); - if i == max_items - 2 { + for (i, item) in self.iter().enumerate() { + if i == max_items.saturating_sub(1) { result.push_str("… "); write!(result, "{}", self.get(self.len() - 1).unwrap()).unwrap(); break; + } else { + write!(result, "{item}").unwrap(); + result.push_str(", "); } } result.push(']'); @@ -1136,7 +1146,7 @@ mod test { ListPrimitiveChunkedBuilder::::new("a", 10, 10, DataType::Int32); builder.append_opt_slice(Some(&[1, 2, 3, 4, 5, 6])); builder.append_opt_slice(None); - let list = builder.finish().into_series(); + let list_long = builder.finish().into_series(); assert_eq!( r#"shape: (2,) @@ -1145,7 +1155,7 @@ Series: 'a' [list[i32]] [1, 2, … 6] null ]"#, - format!("{:?}", list) + format!("{:?}", list_long) ); std::env::set_var("POLARS_FMT_TABLE_CELL_LIST_LEN", "10"); @@ -1157,8 +1167,116 @@ Series: 'a' [list[i32]] [1, 2, 3, 4, 5, 6] null ]"#, - format!("{:?}", list) - ) + format!("{:?}", list_long) + ); + + std::env::set_var("POLARS_FMT_TABLE_CELL_LIST_LEN", "-1"); + + assert_eq!( + r#"shape: (2,) +Series: 'a' [list[i32]] +[ + [1, 2, 3, 4, 5, 6] + null +]"#, + format!("{:?}", list_long) + ); + + std::env::set_var("POLARS_FMT_TABLE_CELL_LIST_LEN", "0"); + + assert_eq!( + r#"shape: (2,) +Series: 'a' [list[i32]] +[ + […] + null +]"#, + format!("{:?}", list_long) + ); + + std::env::set_var("POLARS_FMT_TABLE_CELL_LIST_LEN", "1"); + + assert_eq!( + r#"shape: (2,) +Series: 'a' [list[i32]] +[ + [… 6] + null +]"#, + format!("{:?}", list_long) + ); + + std::env::set_var("POLARS_FMT_TABLE_CELL_LIST_LEN", "4"); + + assert_eq!( + r#"shape: (2,) +Series: 'a' [list[i32]] +[ + [1, 2, 3, … 6] + null +]"#, + format!("{:?}", list_long) + ); + + let mut builder = + ListPrimitiveChunkedBuilder::::new("a", 10, 10, DataType::Int32); + builder.append_opt_slice(Some(&[1])); + builder.append_opt_slice(None); + let list_short = builder.finish().into_series(); + + std::env::set_var("POLARS_FMT_TABLE_CELL_LIST_LEN", ""); + + assert_eq!( + r#"shape: (2,) +Series: 'a' [list[i32]] +[ + [1] + null +]"#, + format!("{:?}", list_short) + ); + + std::env::set_var("POLARS_FMT_TABLE_CELL_LIST_LEN", "0"); + + assert_eq!( + r#"shape: (2,) +Series: 'a' [list[i32]] +[ + […] + null +]"#, + format!("{:?}", list_short) + ); + + std::env::set_var("POLARS_FMT_TABLE_CELL_LIST_LEN", "-1"); + + assert_eq!( + r#"shape: (2,) +Series: 'a' [list[i32]] +[ + [1] + null +]"#, + format!("{:?}", list_short) + ); + + let mut builder = + ListPrimitiveChunkedBuilder::::new("a", 10, 10, DataType::Int32); + builder.append_opt_slice(Some(&[])); + builder.append_opt_slice(None); + let list_empty = builder.finish().into_series(); + + std::env::set_var("POLARS_FMT_TABLE_CELL_LIST_LEN", ""); + + assert_eq!( + r#"shape: (2,) +Series: 'a' [list[i32]] +[ + [] + null +]"#, + format!("{:?}", list_empty) + ); } #[test] diff --git a/py-polars/polars/config.py b/py-polars/polars/config.py index 2d557d914354..3a5a1f1b07cb 100644 --- a/py-polars/polars/config.py +++ b/py-polars/polars/config.py @@ -588,7 +588,9 @@ def set_fmt_table_cell_list_len(cls, n: int | None) -> type[Config]: """ Set the number of elements to display for List values. - Values less than 0 will result in all values being printed. + Empty lists will always print "[]". Negative values will result in all values + being printed. A value of 0 will always "[…]" for lists with contents. A value + of 1 will print only the final item in the list. Parameters ----------