refactor ptable_heatmap()'s tick_fmt() and add test for cbar_precision kwarg

Author: janosh

pymatviz/ptable.py

@@ -302,12 +302,10 @@ def ptable_heatmap(
         def tick_fmt(val: float, _pos: int) -> str:
             # val: value at color axis tick (e.g. 10.0, 20.0, ...)
             # pos: zero-based tick counter (e.g. 0, 1, 2, ...)
-            if heat_mode == "percent":
-                # display color bar values as percentages
-                return f"{val:.0%}"
-            if val < 1e4:
-                return f"{val:{cbar_precision or precision or '.0f'}}"
-            return f"{val:{cbar_precision or precision or '.2g'}}"
+            default = (
+                ".0%" if heat_mode == "percent" else (".0f" if val < 1e4 else ".2g")
+            )
+            return f"{val:{cbar_precision or precision or default}}"
 
         cbar = fig.colorbar(
             mappable, cax=cb_ax, orientation="horizontal", format=tick_fmt

tests/test_ptable.py

@@ -149,6 +149,15 @@ def test_ptable_heatmap(
     with pytest.raises(ValueError, match=r"Unexpected symbol\(s\) foobar"):
         ptable_heatmap(glass_elem_counts, exclude_elements=["foobar"])
 
+    # test cbar_precision
+    ax = ptable_heatmap(glass_elem_counts, cbar_precision=".3f")
+    cbar_1st_label = ax.child_axes[0].get_xticklabels()[0].get_text()
+    assert cbar_1st_label == "0.000"
+
+    ax = ptable_heatmap(glass_elem_counts, heat_mode="percent", cbar_precision=".3%")
+    cbar_1st_label = ax.child_axes[0].get_xticklabels()[0].get_text()
+    assert cbar_1st_label == "0.000%"
+
 
 def test_ptable_heatmap_ratio(
     steel_formulas: list[str],