feat(python, rust): add a rounded_corners modifier to `pl.Config.se…

…t_tbl_formatting` (#6108)
pola-rs · Jan 8, 2023 · aef8938 · aef8938
1 parent 345156d
commit aef8938
Show file tree

Hide file tree

Showing 6 changed files with 92 additions and 55 deletions.
diff --git a/polars/polars-core/src/config.rs b/polars/polars-core/src/config.rs
@@ -12,3 +12,4 @@ pub(crate) const FMT_TABLE_HIDE_DATAFRAME_SHAPE_INFORMATION: &str =
     "POLARS_FMT_TABLE_HIDE_DATAFRAME_SHAPE_INFORMATION";
 pub(crate) const FMT_TABLE_INLINE_COLUMN_DATA_TYPE: &str =
     "POLARS_FMT_TABLE_INLINE_COLUMN_DATA_TYPE";
+pub(crate) const FMT_TABLE_ROUNDED_CORNERS: &str = "POLARS_FMT_TABLE_ROUNDED_CORNERS";
diff --git a/polars/polars-core/src/fmt.rs b/polars/polars-core/src/fmt.rs
@@ -12,6 +12,8 @@ use arrow::temporal_conversions::*;
 #[cfg(feature = "timezones")]
 use chrono::TimeZone;
 #[cfg(any(feature = "fmt", feature = "fmt_no_tty"))]
+use comfy_table::modifiers::*;
+#[cfg(any(feature = "fmt", feature = "fmt_no_tty"))]
 use comfy_table::presets::*;
 #[cfg(any(feature = "fmt", feature = "fmt_no_tty"))]
 use comfy_table::*;
@@ -421,32 +423,35 @@ impl Display for DataFrame {
                 names.push(s);
                 constraints.push(tbl_lower_bounds(l));
             }
-            let preset = match std::env::var(FMT_TABLE_FORMATTING)
+            let (preset, is_utf8) = match std::env::var(FMT_TABLE_FORMATTING)
                 .as_deref()
                 .unwrap_or("DEFAULT")
             {
-                "ASCII_FULL" => ASCII_FULL,
-                "ASCII_FULL_CONDENSED" => ASCII_FULL_CONDENSED,
-                "ASCII_NO_BORDERS" => ASCII_NO_BORDERS,
-                "ASCII_BORDERS_ONLY" => ASCII_BORDERS_ONLY,
-                "ASCII_BORDERS_ONLY_CONDENSED" => ASCII_BORDERS_ONLY_CONDENSED,
-                "ASCII_HORIZONTAL_ONLY" => ASCII_HORIZONTAL_ONLY,
-                "ASCII_MARKDOWN" => ASCII_MARKDOWN,
-                "UTF8_FULL" => UTF8_FULL,
-                "UTF8_FULL_CONDENSED" => UTF8_FULL_CONDENSED,
-                "UTF8_NO_BORDERS" => UTF8_NO_BORDERS,
-                "UTF8_BORDERS_ONLY" => UTF8_BORDERS_ONLY,
-                "UTF8_HORIZONTAL_ONLY" => UTF8_HORIZONTAL_ONLY,
-                "NOTHING" => NOTHING,
-                "DEFAULT" => UTF8_FULL_CONDENSED,
-                _ => UTF8_FULL_CONDENSED,
+                "ASCII_FULL" => (ASCII_FULL, false),
+                "ASCII_FULL_CONDENSED" => (ASCII_FULL_CONDENSED, false),
+                "ASCII_NO_BORDERS" => (ASCII_NO_BORDERS, false),
+                "ASCII_BORDERS_ONLY" => (ASCII_BORDERS_ONLY, false),
+                "ASCII_BORDERS_ONLY_CONDENSED" => (ASCII_BORDERS_ONLY_CONDENSED, false),
+                "ASCII_HORIZONTAL_ONLY" => (ASCII_HORIZONTAL_ONLY, false),
+                "ASCII_MARKDOWN" => (ASCII_MARKDOWN, false),
+                "UTF8_FULL" => (UTF8_FULL, true),
+                "UTF8_FULL_CONDENSED" => (UTF8_FULL_CONDENSED, true),
+                "UTF8_NO_BORDERS" => (UTF8_NO_BORDERS, true),
+                "UTF8_BORDERS_ONLY" => (UTF8_BORDERS_ONLY, true),
+                "UTF8_HORIZONTAL_ONLY" => (UTF8_HORIZONTAL_ONLY, true),
+                "NOTHING" => (NOTHING, false),
+                "DEFAULT" => (UTF8_FULL_CONDENSED, true),
+                _ => (UTF8_FULL_CONDENSED, true),
             };
 
             let mut table = Table::new();
             table
                 .load_preset(preset)
                 .set_content_arrangement(ContentArrangement::Dynamic);
 
+            if is_utf8 && env_is_true(FMT_TABLE_ROUNDED_CORNERS) {
+                table.apply_modifier(UTF8_ROUND_CORNERS);
+            }
             if max_n_rows > 0 {
                 if height > max_n_rows {
                     let mut rows = Vec::with_capacity(std::cmp::max(max_n_rows, 2));

diff --git a/polars/src/lib.rs b/polars/src/lib.rs
@@ -343,28 +343,31 @@
 //!                                    LEFT
 //!                                    CENTER
 //!                                    RIGHT
-//! * `POLARS_FMT_TABLE_HIDE_COLUMN_NAMES` -> Hide column names
-//! * `POLARS_FMT_TABLE_HIDE_COLUMN_DATA_TYPES` -> Hide data types for columns
-//! * `POLARS_FMT_TABLE_HIDE_COLUMN_SEPARATOR` -> Hide separator that separates column names from rows
-//! * `POLARS_FMT_TABLE_CHANGE_COLUMN_DATA_TYPE_POSITION_FORMAT"` -> changes the position and format of data, putting it directly below column name
+//! * `POLARS_FMT_TABLE_DATAFRAME_SHAPE_BELOW` -> print shape information below the table.
+//! * `POLARS_FMT_TABLE_HIDE_COLUMN_NAMES` -> hide table column names.
+//! * `POLARS_FMT_TABLE_HIDE_COLUMN_DATA_TYPES` -> hide data types for columns.
+//! * `POLARS_FMT_TABLE_HIDE_COLUMN_SEPARATOR` -> hide separator that separates column names from rows.
+//! * `POLARS_FMT_TABLE_HIDE_DATAFRAME_SHAPE_INFORMATION"` -> omit table shape information.
+//! * `POLARS_FMT_TABLE_INLINE_COLUMN_DATA_TYPE` -> put column data type on the same line as the column name.
+//! * `POLARS_FMT_TABLE_ROUNDED_CORNERS` -> apply rounded corners to UTF8-styled tables.
 //! * `POLARS_FMT_MAX_COLS` -> maximum number of columns shown when formatting DataFrames.
 //! * `POLARS_FMT_MAX_ROWS` -> maximum number of rows shown when formatting DataFrames.
 //! * `POLARS_FMT_STR_LEN` -> maximum number of characters printed per string value.
 //! * `POLARS_TABLE_WIDTH` -> width of the tables used during DataFrame formatting.
 //! * `POLARS_MAX_THREADS` -> maximum number of threads used to initialize thread pool (on startup).
-//! * `POLARS_VERBOSE` -> print logging info to stderr
-//! * `POLARS_NO_PARTITION` -> Polars may choose to partition the groupby operation, based on data
-//!                            cardinality. Setting this env var will turn partitioned groupby's off
+//! * `POLARS_VERBOSE` -> print logging info to stderr.
+//! * `POLARS_NO_PARTITION` -> polars may choose to partition the groupby operation, based on data
+//!                            cardinality. Setting this env var will turn partitioned groupby's off.
 //! * `POLARS_PARTITION_SAMPLE_FRAC` -> how large chunk of the dataset to sample to determine cardinality,
-//!                                     defaults to `0.001`
+//!                                     defaults to `0.001`.
 //! * `POLARS_PARTITION_UNIQUE_COUNT` -> at which (estimated) key count a partitioned groupby should run.
 //!                                          defaults to `1000`, any higher cardinality will run default groupby.
-//! * `POLARS_FORCE_PARTITION` -> Force partitioned groupby if the keys and aggregations allow it.
+//! * `POLARS_FORCE_PARTITION` -> force partitioned groupby if the keys and aggregations allow it.
 //! * `POLARS_ALLOW_EXTENSION` -> allows for `[ObjectChunked<T>]` to be used in arrow, opening up possibilities like using
 //!                               `T` in complex lazy expressions. However this does require `unsafe` code allow this.
 //! * `POLARS_NO_PARQUET_STATISTICS` -> if set, statistics in parquet files are ignored.
 //! * `POLARS_PANIC_ON_ERR` -> panic instead of returning an Error.
-//! * `POLARS_NO_CHUNKED_JOIN` -> Force rechunk before joins.
+//! * `POLARS_NO_CHUNKED_JOIN` -> force rechunk before joins.
 //!
 //!
 //! ## User Guide

diff --git a/py-polars/polars/cfg.py b/py-polars/polars/cfg.py
@@ -26,6 +26,7 @@
     "POLARS_FMT_TABLE_HIDE_COLUMN_SEPARATOR",
     "POLARS_FMT_TABLE_HIDE_DATAFRAME_SHAPE_INFORMATION",
     "POLARS_FMT_TABLE_INLINE_COLUMN_DATA_TYPE",
+    "POLARS_FMT_TABLE_ROUNDED_CORNERS",
     "POLARS_TABLE_WIDTH",
     "POLARS_VERBOSE",
 }
@@ -328,21 +329,25 @@ def set_tbl_dataframe_shape_below(cls, active: bool = True) -> type[Config]:
     @classmethod
     def set_tbl_formatting(
         cls,
-        format: Literal[
-            "ASCII_FULL",
-            "ASCII_FULL_CONDENSED",
-            "ASCII_NO_BORDERS",
-            "ASCII_BORDERS_ONLY",
-            "ASCII_BORDERS_ONLY_CONDENSED",
-            "ASCII_HORIZONTAL_ONLY",
-            "ASCII_MARKDOWN",
-            "UTF8_FULL",
-            "UTF8_FULL_CONDENSED",
-            "UTF8_NO_BORDERS",
-            "UTF8_BORDERS_ONLY",
-            "UTF8_HORIZONTAL_ONLY",
-            "NOTHING",
-        ],
+        format: (
+            Literal[
+                "ASCII_FULL",
+                "ASCII_FULL_CONDENSED",
+                "ASCII_NO_BORDERS",
+                "ASCII_BORDERS_ONLY",
+                "ASCII_BORDERS_ONLY_CONDENSED",
+                "ASCII_HORIZONTAL_ONLY",
+                "ASCII_MARKDOWN",
+                "UTF8_FULL",
+                "UTF8_FULL_CONDENSED",
+                "UTF8_NO_BORDERS",
+                "UTF8_BORDERS_ONLY",
+                "UTF8_HORIZONTAL_ONLY",
+                "NOTHING",
+            ]
+            | None
+        ) = None,
+        rounded_corners: bool = False,
     ) -> type[Config]:
         """
         Set table formatting style.
@@ -364,6 +369,9 @@ def set_tbl_formatting(
             * "UTF8_HORIZONTAL_ONLY": UTF8, horizontal lines only.
             * "NOTHING": No borders or other lines.
 
+        rounded_corners : bool
+            apply rounded corners to UTF8-styled tables (no-op for ASCII formats).
+
         Notes
         -----
         The UTF8 styles all use one or more of the semigraphic box-drawing characters
@@ -377,7 +385,9 @@ def set_tbl_formatting(
         """
         # can see what the different styles look like in the comfy-table tests:
         # https://github.com/Nukesor/comfy-table/blob/main/tests/all/presets_test.rs
-        os.environ["POLARS_FMT_TABLE_FORMATTING"] = format
+        if format:
+            os.environ["POLARS_FMT_TABLE_FORMATTING"] = format
+        os.environ["POLARS_FMT_TABLE_ROUNDED_CORNERS"] = str(int(rounded_corners))
         return cls
 
     @classmethod

diff --git a/py-polars/polars/internals/lazy_functions.py b/py-polars/polars/internals/lazy_functions.py
@@ -2455,10 +2455,20 @@ def arg_where(
 
 
 def coalesce(
-    exprs: Sequence[
-        pli.Expr | str | date | datetime | timedelta | int | float | bool | pli.Series
-    ]
-    | pli.Expr,
+    exprs: (
+        Sequence[
+            pli.Expr
+            | str
+            | date
+            | datetime
+            | timedelta
+            | int
+            | float
+            | bool
+            | pli.Series
+        ]
+        | pli.Expr
+    ),
 ) -> pli.Expr:
     """
     Folds the expressions from left to right, keeping the first non-null value.

diff --git a/py-polars/tests/unit/test_cfg.py b/py-polars/tests/unit/test_cfg.py
@@ -344,32 +344,35 @@ def test_shape_below_table_and_inlined_dtype() -> None:
     df = pl.DataFrame({"a": [1, 2], "b": [3, 4], "c": [5, 6]})
 
     pl.Config.set_tbl_column_data_type_inline(True).set_tbl_dataframe_shape_below(True)
-    pl.Config.set_tbl_formatting("UTF8_FULL")
+    pl.Config.set_tbl_formatting("UTF8_FULL", rounded_corners=True)
     assert (
         str(df) == ""
-        "┌─────────┬─────────┬─────────┐\n"
+        "╭─────────┬─────────┬─────────╮\n"
         "│ a (i64) ┆ b (i64) ┆ c (i64) │\n"
         "╞═════════╪═════════╪═════════╡\n"
         "│ 1       ┆ 3       ┆ 5       │\n"
         "├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤\n"
         "│ 2       ┆ 4       ┆ 6       │\n"
-        "└─────────┴─────────┴─────────┘\n"
+        "╰─────────┴─────────┴─────────╯\n"
         "shape: (2, 3)"
     )
 
     pl.Config.set_tbl_dataframe_shape_below(False)
     assert (
         str(df) == "shape: (2, 3)\n"
-        "┌─────────┬─────────┬─────────┐\n"
+        "╭─────────┬─────────┬─────────╮\n"
         "│ a (i64) ┆ b (i64) ┆ c (i64) │\n"
         "╞═════════╪═════════╪═════════╡\n"
         "│ 1       ┆ 3       ┆ 5       │\n"
         "├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤\n"
         "│ 2       ┆ 4       ┆ 6       │\n"
-        "└─────────┴─────────┴─────────┘"
+        "╰─────────┴─────────┴─────────╯"
+    )
+    (
+        pl.Config.set_tbl_formatting(rounded_corners=False)
+        .set_tbl_column_data_type_inline(False)
+        .set_tbl_cell_alignment("RIGHT")
     )
-
-    pl.Config.set_tbl_column_data_type_inline(False).set_tbl_cell_alignment("RIGHT")
     assert (
         str(df) == "shape: (2, 3)\n"
         "┌─────┬─────┬─────┐\n"
@@ -444,13 +447,18 @@ def test_config_scope() -> None:
     initial_state = pl.Config.state()
 
     with pl.Config() as cfg:
-        cfg.set_verbose(True).set_tbl_hide_dtype_separator(True).set_ascii_tables()
-
+        (
+            cfg.set_tbl_formatting(rounded_corners=True)
+            .set_verbose(True)
+            .set_tbl_hide_dtype_separator(True)
+            .set_ascii_tables()
+        )
         new_state_entries = set(
             {
                 "POLARS_FMT_MAX_COLS": "8",
                 "POLARS_FMT_TABLE_FORMATTING": "ASCII_FULL_CONDENSED",
                 "POLARS_FMT_TABLE_HIDE_COLUMN_SEPARATOR": "1",
+                "POLARS_FMT_TABLE_ROUNDED_CORNERS": "1",
                 "POLARS_VERBOSE": "1",
             }.items()
         )