Coverage for src / rtflite / pagination / strategies / defaults.py: 96%

1from typing import cast 

2 

3from ..core import PageBreakCalculator, RTFPagination 

4from .base import PageContext, PaginationContext, PaginationStrategy 

5 

6 

7class DefaultPaginationStrategy(PaginationStrategy): 

8 """Default pagination strategy based on row counts and page size.""" 

9 

10 def paginate(self, context: PaginationContext) -> list[PageContext]: 

11 # Initialize calculator 

12 assert context.rtf_page.width is not None 

13 assert context.rtf_page.height is not None 

14 assert context.rtf_page.margin is not None 

15 assert context.rtf_page.nrow is not None 

16 assert context.rtf_page.orientation is not None 

17 

18 pagination_config = RTFPagination( 

19 page_width=context.rtf_page.width, 

20 page_height=context.rtf_page.height, 

21 margin=context.rtf_page.margin, 

22 nrow=context.rtf_page.nrow, 

23 orientation=context.rtf_page.orientation, 

24 ) 

25 calculator = PageBreakCalculator(pagination=pagination_config) 

26 

27 # Calculate metadata 

28 metadata = calculator.calculate_row_metadata( 

29 df=context.df, 

30 col_widths=context.col_widths, 

31 table_attrs=context.table_attrs, 

32 removed_column_indices=context.removed_column_indices, 

33 additional_rows_per_page=context.additional_rows_per_page, 

34 ) 

35 

36 # Create PageContext objects 

37 pages = [] 

38 import polars as pl 

39 

40 # Get unique pages and sort them 

41 unique_pages = metadata["page"].unique().sort() 

42 total_pages = len(unique_pages) 

43 

44 for page_num in unique_pages: 

45 # Filter metadata for this page 

46 page_rows = metadata.filter(pl.col("page") == page_num) 

47 

48 if page_rows.height == 0: 

49 continue 

50 

51 start_row = cast(int, page_rows["row_index"].min()) 

52 end_row = cast(int, page_rows["row_index"].max()) 

53 

54 # Slice the original dataframe 

55 # Note: end_row is inclusive index, slice takes length 

56 page_df = context.df.slice(start_row, end_row - start_row + 1) 

57 

58 # 1-based page number for display 

59 display_page_num = int(page_num) 

60 

61 pages.append( 

62 PageContext( 

63 page_number=display_page_num, 

64 total_pages=total_pages, 

65 data=page_df, 

66 is_first_page=(display_page_num == 1), 

67 is_last_page=(display_page_num == total_pages), 

68 col_widths=context.col_widths, 

69 needs_header=( 

70 context.rtf_body.pageby_header or display_page_num == 1 

71 ), 

72 table_attrs=context.table_attrs, 

73 ) 

74 ) 

75 

76 return pages