rebase and add 'unset' logic for Config.set_tbl_cell_numeric_alignment

pola-rs · Sep 15, 2023 · 0abc339 · 0abc339
1 parent dca1ea6
commit 0abc339
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/py-polars/polars/config.py b/py-polars/polars/config.py
@@ -528,7 +528,7 @@ def set_tbl_cell_alignment(
 
     @classmethod
     def set_tbl_cell_numeric_alignment(
-        cls, format: Literal["LEFT", "CENTER", "RIGHT"]
+        cls, format: Literal["LEFT", "CENTER", "RIGHT"] | None
     ) -> type[Config]:
         """
         Set table cell alignment for numeric columns.
@@ -568,7 +568,12 @@ def set_tbl_cell_numeric_alignment(
         KeyError: if alignment string not recognised.
 
         """
-        os.environ["POLARS_FMT_TABLE_CELL_NUMERIC_ALIGNMENT"] = format
+        if format is None:
+            os.environ.pop("POLARS_FMT_TABLE_CELL_NUMERIC_ALIGNMENT", None)
+        elif format not in {"LEFT", "CENTER", "RIGHT"}:
+            raise ValueError(f"invalid alignment: {format!r}")
+        else:
+            os.environ["POLARS_FMT_TABLE_CELL_NUMERIC_ALIGNMENT"] = format
         return cls
 
     @classmethod

diff --git a/py-polars/tests/unit/test_cfg.py b/py-polars/tests/unit/test_cfg.py
@@ -766,6 +766,12 @@ def test_set_fmt_str_lengths_invalid_length() -> None:
         ("POLARS_FMT_MAX_ROWS", "set_tbl_rows", 3, "3"),
         ("POLARS_FMT_STR_LEN", "set_fmt_str_lengths", 42, "42"),
         ("POLARS_FMT_TABLE_CELL_ALIGNMENT", "set_tbl_cell_alignment", "RIGHT", "RIGHT"),
+        (
+            "POLARS_FMT_TABLE_CELL_NUMERIC_ALIGNMENT",
+            "set_tbl_cell_numeric_alignment",
+            "RIGHT",
+            "RIGHT",
+        ),
         ("POLARS_FMT_TABLE_HIDE_COLUMN_NAMES", "set_tbl_hide_column_names", True, "1"),
         (
             "POLARS_FMT_TABLE_DATAFRAME_SHAPE_BELOW",