Coverage for barbet/output.py: 81.25%

1from typing import TYPE_CHECKING 

2 

3if TYPE_CHECKING: 

4 import polars as pl 

5 

6 

7def print_polars_df(df: "pl.DataFrame", max_rows: int = 20, column_names:list[str]|None=None): 

8 from rich.table import Table 

9 from rich.console import Console 

10 

11 console = Console() 

12 table = Table(show_header=True, header_style="bold magenta") 

13 

14 # Add columns 

15 for col in column_names or df.columns: 

16 table.add_column(col) 

17 

18 # Add rows 

19 for row in df.head(max_rows).iter_rows(): 

20 table.add_row(*[str(cell) for cell in row]) 

21 

22 # Indicate truncation if needed 

23 if df.height > max_rows: 

24 table.add_row(*["..." for _ in df.columns]) 

25 console.print(f"[dim]Showing first {max_rows} of {df.height} rows[/dim]") 

26 

27 console.print(table)