Matplotlib table: use bold weight when calculating header width (#90)

* use bold for header width * add test for long column headers
dexplo · May 14, 2023 · f4787e1 · f4787e1
1 parent eff1f4d
commit f4787e1
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 4 deletions.
diff --git a/dataframe_image/_matplotlib_table.py b/dataframe_image/_matplotlib_table.py
@@ -147,20 +147,23 @@ def parse_row(row):
                 rows.append(parse_row(row))
         return rows, num_header_rows
 
-    def get_text_width(self, text):
+    def get_text_width(self, text, weight=None):
         fig = self.text_fig
-        t = fig.text(0, 0, text, size=self.fontsize)
+        t = fig.text(0, 0, text, size=self.fontsize, weight=weight)
         bbox = t.get_window_extent(renderer=self.renderer)
         return bbox.width
 
     def get_all_text_widths(self, rows):
         all_text_widths = []
-        for row in rows:
+        for i, row in enumerate(rows):
             row_widths = []
             for vals in row:
                 cell_max_width = 0
                 for text in vals[0].split("\n"):
-                    cell_max_width = max(cell_max_width, self.get_text_width(text))
+                    weight = "bold" if i == 0 else None
+                    cell_max_width = max(
+                        cell_max_width, self.get_text_width(text, weight)
+                    )
                 row_widths.append(cell_max_width)
             all_text_widths.append(row_widths)
         pad = 10  # number of pixels to pad columns with

diff --git a/tests/test_df_image.py b/tests/test_df_image.py
@@ -1,6 +1,8 @@
 import numpy as np
 import pandas as pd
 import pytest
+import random
+import string
 
 import dataframe_image
 
@@ -41,3 +43,24 @@ def test_huge_df(self, converter, dpi):
             dpi=dpi,
             max_rows=-1,
         )
+
+    @pytest.mark.parametrize("dpi", test_dpi_values)
+    @pytest.mark.parametrize("converter", converters)
+    def test_long_column_headers(self, converter, dpi):
+        column_headers = [
+            "".join(
+                random.choices(string.ascii_uppercase + string.ascii_lowercase, k=80)
+            )
+            for _ in range(5)
+        ]
+
+        df = pd.DataFrame(
+            np.random.randint(0, 100, size=(5, 5)), columns=column_headers
+        )
+
+        df.dfi.export(
+            f"tests/test_output/long_column_{converter}_dpi{dpi}.png",
+            table_conversion=converter,
+            dpi=dpi,
+            max_rows=-1,
+        )