diff --git a/dataframe_image/_matplotlib_table.py b/dataframe_image/_matplotlib_table.py index 125b2ca..f155839 100644 --- 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 index b68d0ad..90d7201 100644 --- 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, + )