Fix fmt when set_tbl_rows/cols (#2230)

.gitignore

@@ -11,4 +11,5 @@ polars/vendor
 .env
 AUTO_CHANGELOG.md
 node_modules/
-.coverage
+.coverage
+venv/

polars/polars-core/src/fmt.rs

@@ -399,16 +399,16 @@ impl Display for DataFrame {
                 }
                 let dots = rows[0].iter().map(|_| "...".to_string()).collect();
                 rows.push(dots);
-                for i in (self.height() - max_n_rows / 2 - 1)..self.height() {
+                for i in (self.height() - (max_n_rows + 1) / 2)..self.height() {
                     let row = self.columns.iter().map(|s| s.str_value(i)).collect();
                     rows.push(prepare_row(row, n_first, n_last));
                 }
                 for row in rows {
                     table.add_row(row);
                 }
             } else {
-                for i in 0..max_n_rows {
-                    if i < self.height() && self.width() > 0 {
+                for i in 0..self.height() {
+                    if self.width() > 0 {
                         let row = self.columns.iter().map(|s| s.str_value(i)).collect();
                         table.add_row(prepare_row(row, n_first, n_last));
                     } else {
@@ -440,16 +440,16 @@ impl Display for DataFrame {
                 }
                 let dots = rows[0].iter().map(|_| "...".to_string()).collect();
                 rows.push(dots);
-                for i in (self.height() - max_n_rows / 2 - 1)..self.height() {
+                for i in (self.height() - (max_n_rows + 1) / 2)..self.height() {
                     let row = self.columns.iter().map(|s| s.str_value(i)).collect();
                     rows.push(prepare_row(row, n_first, n_last));
                 }
                 for row in rows {
                     table.add_row(Row::new(row.into_iter().map(|s| Cell::new(&s)).collect()));
                 }
             } else {
-                for i in 0..max_n_rows {
-                    if i < self.height() && self.width() > 0 {
+                for i in 0..self.height() {
+                    if self.width() > 0 {
                         let row = self.columns.iter().map(|s| s.str_value(i)).collect();
                         table.add_row(Row::new(
                             prepare_row(row, n_first, n_last)

py-polars/polars/internals/expr.py

@@ -1050,8 +1050,6 @@ def over(self, expr: Union[str, "Expr", List[Union["Expr", str]]]) -> "Expr":
         ├╌╌╌╌╌╌╌╌┤
         │ ...    │
         ├╌╌╌╌╌╌╌╌┤
-        │ 4      │
-        ├╌╌╌╌╌╌╌╌┤
         │ 6      │
         ├╌╌╌╌╌╌╌╌┤
         │ 6      │

py-polars/polars/internals/frame.py

@@ -3122,8 +3122,6 @@ def explode(
         ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
         │ ...     ┆ ... │
         ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
-        │ c       ┆ 5   │
-        ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
         │ a       ┆ 6   │
         ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
         │ b       ┆ 2   │

py-polars/polars/internals/lazy_frame.py

@@ -1024,8 +1024,6 @@ def explode(
         ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
         │ ...     ┆ ... │
         ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
-        │ c       ┆ 5   │
-        ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
         │ a       ┆ 6   │
         ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
         │ b       ┆ 2   │

py-polars/tests/test_cfg.py

@@ -66,7 +66,7 @@ def test_tbl_width_chars(environ: None) -> None:
     assert max(len(line) for line in str(df).split("\n")) == 13
 
 
-def test_tbl_cols(environ: None) -> None:
+def test_set_tbl_cols(environ: None) -> None:
     df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})
 
     pl.Config.set_tbl_cols(1)
@@ -76,13 +76,108 @@ def test_tbl_cols(environ: None) -> None:
     pl.Config.set_tbl_cols(3)
     assert str(df).split("\n")[2] == "│ a   ┆ b   ┆ c   │"
 
+    df = pl.DataFrame(
+        {"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9], "d": [10, 11, 12]}
+    )
+    pl.Config.set_tbl_cols(2)
+    assert str(df).split("\n")[2] == "│ a   ┆ ... ┆ d   │"
+    pl.Config.set_tbl_cols(3)
+    assert str(df).split("\n")[2] == "│ a   ┆ b   ┆ ... ┆ d   │"
+
 
-@pytest.mark.skip("not correctly implemented at the moment")
 def test_set_tbl_rows(environ: None) -> None:
-    # df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})
-    pass
-    #
-    # pl.Config.set_tbl_rows(3)
+    df = pl.DataFrame({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8], "c": [9, 10, 11, 12]})
+
+    pl.Config.set_tbl_rows(1)
+    assert (
+        str(df) == "shape: (4, 3)\n"
+        "┌─────┬─────┬─────┐\n"
+        "│ a   ┆ b   ┆ c   │\n"
+        "│ --- ┆ --- ┆ --- │\n"
+        "│ i64 ┆ i64 ┆ i64 │\n"
+        "╞═════╪═════╪═════╡\n"
+        "│ 1   ┆ 5   ┆ 9   │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ ... ┆ ... ┆ ... │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ 4   ┆ 8   ┆ 12  │\n"
+        "└─────┴─────┴─────┘"
+    )
+    pl.Config.set_tbl_rows(2)
+    assert (
+        str(df) == "shape: (4, 3)\n"
+        "┌─────┬─────┬─────┐\n"
+        "│ a   ┆ b   ┆ c   │\n"
+        "│ --- ┆ --- ┆ --- │\n"
+        "│ i64 ┆ i64 ┆ i64 │\n"
+        "╞═════╪═════╪═════╡\n"
+        "│ 1   ┆ 5   ┆ 9   │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ ... ┆ ... ┆ ... │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ 4   ┆ 8   ┆ 12  │\n"
+        "└─────┴─────┴─────┘"
+    )
+    pl.Config.set_tbl_rows(3)
+    assert (
+        str(df) == "shape: (4, 3)\n"
+        "┌─────┬─────┬─────┐\n"
+        "│ a   ┆ b   ┆ c   │\n"
+        "│ --- ┆ --- ┆ --- │\n"
+        "│ i64 ┆ i64 ┆ i64 │\n"
+        "╞═════╪═════╪═════╡\n"
+        "│ 1   ┆ 5   ┆ 9   │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ ... ┆ ... ┆ ... │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ 3   ┆ 7   ┆ 11  │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ 4   ┆ 8   ┆ 12  │\n"
+        "└─────┴─────┴─────┘"
+    )
+    pl.Config.set_tbl_rows(4)
+    assert (
+        str(df) == "shape: (4, 3)\n"
+        "┌─────┬─────┬─────┐\n"
+        "│ a   ┆ b   ┆ c   │\n"
+        "│ --- ┆ --- ┆ --- │\n"
+        "│ i64 ┆ i64 ┆ i64 │\n"
+        "╞═════╪═════╪═════╡\n"
+        "│ 1   ┆ 5   ┆ 9   │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ 2   ┆ 6   ┆ 10  │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ 3   ┆ 7   ┆ 11  │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ 4   ┆ 8   ┆ 12  │\n"
+        "└─────┴─────┴─────┘"
+    )
+
+    df = pl.DataFrame(
+        {
+            "a": [1, 2, 3, 4, 5],
+            "b": [6, 7, 8, 9, 10],
+            "c": [11, 12, 13, 14, 15],
+        }
+    )
+
+    pl.Config.set_tbl_rows(3)
+    assert (
+        str(df) == "shape: (5, 3)\n"
+        "┌─────┬─────┬─────┐\n"
+        "│ a   ┆ b   ┆ c   │\n"
+        "│ --- ┆ --- ┆ --- │\n"
+        "│ i64 ┆ i64 ┆ i64 │\n"
+        "╞═════╪═════╪═════╡\n"
+        "│ 1   ┆ 6   ┆ 11  │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ ... ┆ ... ┆ ... │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ 4   ┆ 9   ┆ 14  │\n"
+        "├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤\n"
+        "│ 5   ┆ 10  ┆ 15  │\n"
+        "└─────┴─────┴─────┘"
+    )
 
 
 def test_string_cache(environ: None) -> None: